Create train_infer.py

train_infer.py

@@ -0,0 +1,80 @@
+#!/usr/bin/env python3
+running_loss += loss.item() * imgs.size(0)
+
+
+# Validation
+model.eval()
+correct = 0
+total = 0
+with torch.no_grad():
+for imgs, labels in val_loader:
+imgs, labels = imgs.to(device), labels.to(device)
+outputs = model(imgs)
+_, preds = outputs.max(1)
+correct += (preds == labels).sum().item()
+total += labels.size(0)
+val_acc = correct / total if total else 0
+avg_loss = running_loss / (len(train_loader.dataset) if len(train_loader.dataset) else 1)
+print(f"Epoch {epoch+1}/{args.epochs} val_acc={val_acc:.4f} train_loss={avg_loss:.4f}")
+
+
+if val_acc > best_acc:
+best_acc = val_acc
+torch.save({'model_state': model.state_dict(), 'classes': classes, 'base': args.base}, args.ckpt)
+print("Saved best checkpoint ->", args.ckpt)
+
+
+print("Training finished. Best val acc:", best_acc)
+
+
+
+
+def predict(args):
+if not os.path.exists(args.ckpt):
+raise SystemExit(f"Checkpoint not found: {args.ckpt}")
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+ck = torch.load(args.ckpt, map_location=device)
+classes = ck['classes']
+base = ck.get('base', 'resnet18')
+model = build_model(len(classes), base=base, pretrained=False)
+model.load_state_dict(ck['model_state'])
+model.to(device).eval()
+
+
+tf = transforms.Compose([
+transforms.Resize(256),
+transforms.CenterCrop(args.img),
+transforms.ToTensor(),
+transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])
+])
+img = Image.open(args.image).convert('RGB')
+x = tf(img).unsqueeze(0).to(device)
+with torch.no_grad():
+out = model(x)
+pred = out.argmax(1).item()
+print('Predicted:', classes[pred])
+
+
+
+
+if __name__ == '__main__':
+p = argparse.ArgumentParser()
+p.add_argument('--mode', choices=['train', 'predict'], required=True)
+p.add_argument('--train_dir', default='data/train')
+p.add_argument('--val_dir', default='data/val')
+p.add_argument('--image', help='Image path for prediction')
+p.add_argument('--ckpt', default='ckpt.pth')
+p.add_argument('--epochs', type=int, default=3)
+p.add_argument('--batch', type=int, default=16)
+p.add_argument('--lr', type=float, default=1e-4)
+p.add_argument('--img', type=int, default=224)
+p.add_argument('--base', default='resnet18')
+args = p.parse_args()
+
+
+if args.mode == 'train':
+train(args)
+else:
+if not args.image:
+raise SystemExit('Provide --image for predict mode')
+predict(args)