import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torch.cuda.amp import autocast, GradScaler
import torchvision
from torchvision import transforms, datasets, models
from tqdm import tqdm
import numpy as np
config = {
'data_dir': './data',
'num_classes': 100,
'batch_size': 128,
'epochs': 90,
'lr': 0.1,
'momentum': 0.9,
'weight_decay': 1e-4,
'num_workers': 8,
'seed': 42,
}
torch.manual_seed(config['seed'])
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
train_transform = transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ColorJitter(0.2, 0.2, 0.2),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])
val_transform = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])
train_dataset = datasets.ImageFolder(f'{config["data_dir"]}/train', train_transform)
val_dataset = datasets.ImageFolder(f'{config["data_dir"]}/val', val_transform)
train_loader = DataLoader(
train_dataset, batch_size=config['batch_size'], shuffle=True,
num_workers=config['num_workers'], pin_memory=True, drop_last=True
)
val_loader = DataLoader(
val_dataset, batch_size=config['batch_size'], shuffle=False,
num_workers=config['num_workers'], pin_memory=True
)
model = models.resnet50(weights=None, num_classes=config['num_classes'])
model = model.to(device)
criterion = nn.CrossEntropyLoss(label_smoothing=0.1)
optimizer = optim.SGD(
model.parameters(),
lr=config['lr'],
momentum=config['momentum'],
weight_decay=config['weight_decay'],
nesterov=True
)
scheduler = optim.lr_scheduler.CosineAnnealingLR(
optimizer, T_max=config['epochs']
)
scaler = GradScaler()
def train_one_epoch(model, loader, criterion, optimizer, scaler, epoch):
model.train()
running_loss = 0.0
correct = 0
total = 0
pbar = tqdm(loader, desc=f'Epoch {epoch}')
for images, labels in pbar:
images, labels = images.to(device), labels.to(device)
optimizer.zero_grad(set_to_none=True)
with autocast():
outputs = model(images)
loss = criterion(outputs, labels)
scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
scaler.step(optimizer)
scaler.update()
running_loss += loss.item()
_, predicted = outputs.max(1)
correct += predicted.eq(labels).sum().item()
total += labels.size(0)
pbar.set_postfix({
'loss': f'{loss.item():.3f}',
'acc': f'{100.*correct/total:.1f}%'
})
return running_loss / len(loader), 100. * correct / total
@torch.no_grad()
def validate(model, loader, criterion):
model.eval()
running_loss = 0.0
correct = 0
total = 0
for images, labels in tqdm(loader, desc='Validation'):
images, labels = images.to(device), labels.to(device)
outputs = model(images)
loss = criterion(outputs, labels)
running_loss += loss.item()
_, predicted = outputs.max(1)
correct += predicted.eq(labels).sum().item()
total += labels.size(0)
return running_loss / len(loader), 100. * correct / total
best_acc = 0.0
for epoch in range(config['epochs']):
train_loss, train_acc = train_one_epoch(
model, train_loader, criterion, optimizer, scaler, epoch
)
scheduler.step()
val_loss, val_acc = validate(model, val_loader, criterion)
print(f'Epoch {epoch}: '
f'Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.2f}% | '
f'Val Loss: {val_loss:.4f}, Val Acc: {val_acc:.2f}%')
if val_acc > best_acc:
best_acc = val_acc
torch.save({
'epoch': epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'best_acc': best_acc,
}, 'best_model.pth')
print(f' Saved best model with accuracy {best_acc:.2f}%')
print(f'Training completed. Best accuracy: {best_acc:.2f}%')
@torch.no_grad()
def predict(image_path, model, transform, class_names, top_k=5):
from PIL import Image
model.eval()
image = Image.open(image_path).convert('RGB')
input_tensor = transform(image).unsqueeze(0).to(device)
output = model(input_tensor)
probs = torch.nn.functional.softmax(output, dim=1)
top_probs, top_indices = probs.topk(top_k, dim=1)
results = []
for prob, idx in zip(top_probs[0].cpu().numpy(), top_indices[0].cpu().numpy()):
results.append({
'class': class_names[idx],
'probability': float(prob)
})
return results
|
