Spaces:
Sleeping
Sleeping
File size: 3,444 Bytes
26c2a4a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 |
import torch
import torch.nn as nn
import torch.optim as optim
from tqdm import tqdm
import yaml
import os
# Import our project modules
from src.data.data_loader import create_dataloaders
from src.models.model import VisionGuardModel
def train_one_epoch(model, loader, criterion, optimizer, device):
model.train()
running_loss = 0.0
correct = 0
total = 0
loop = tqdm(loader, leave=False)
for images, labels in loop:
images, labels = images.to(device), labels.to(device)
# Forward
outputs = model(images)
loss = criterion(outputs, labels)
# Backward
optimizer.zero_grad()
loss.backward()
optimizer.step()
# Stats
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
running_loss += loss.item()
loop.set_description(f"Loss: {loss.item():.4f}")
return running_loss / len(loader), 100 * correct / total
def validate(model, loader, criterion, device):
model.eval()
correct = 0
total = 0
running_loss = 0.0
with torch.no_grad():
for images, labels in loader:
images, labels = images.to(device), labels.to(device)
outputs = model(images)
loss = criterion(outputs, labels)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
running_loss += loss.item()
return running_loss / len(loader), 100 * correct / total
def main():
# 1. Config & Device
with open("configs/config.yaml", "r") as f:
cfg = yaml.safe_load(f)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"🚀 Training on: {device}")
# 2. Save Path (Google Drive)
save_dir = "/content/drive/MyDrive/VisionGuard_Models"
os.makedirs(save_dir, exist_ok=True)
save_path = f"{save_dir}/dinov2_best.pt"
# 3. Load Data & Model
train_loader, val_loader = create_dataloaders("configs/config.yaml")
model = VisionGuardModel(num_classes=2).to(device)
# 4. Optimizer (Only training the Head)
# Note: We only pass model.head.parameters() to optimizer because backbone is frozen!
optimizer = optim.AdamW(model.head.parameters(), lr=float(cfg['model']['learning_rate_head']))
criterion = nn.CrossEntropyLoss()
# 5. Training Loop
best_acc = 0.0
epochs = cfg['model']['epochs']
print(f"\n🔥 Starting Training for {epochs} Epochs...")
print(f"💾 Best model will be saved to: {save_path}")
for epoch in range(epochs):
print(f"\nEpoch {epoch+1}/{epochs}")
train_loss, train_acc = train_one_epoch(model, train_loader, criterion, optimizer, device)
val_loss, val_acc = validate(model, val_loader, criterion, device)
print(f" Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.2f}%")
print(f" Val Loss: {val_loss:.4f} | Val Acc: {val_acc:.2f}%")
# Save Best
if val_acc > best_acc:
best_acc = val_acc
torch.save(model.state_dict(), save_path)
print(f" ⭐ Saved New Best Model ({best_acc:.2f}%)")
if __name__ == "__main__":
main() |