Image Classification Project in Python: Transfer Learning With PyTorch

How to build an image classifier with PyTorch, step by step

1. Step 1. Set up PyTorch and torchvision

   On your own machine, create an environment and install the libraries. On Colab, torch and torchvision are usually preinstalled.

   mkdir image-classifier && cd image-classifier
   python -m venv .venv
   # Windows: .venv\Scripts\activate    macOS/Linux: source .venv/bin/activate
   pip install torch torchvision scikit-learn matplotlib pillow
   python -c "import torch; print(torch.__version__, torch.cuda.is_available())"

   If the last line prints True, PyTorch can see your GPU. If you need a specific CUDA build, use the install selector on the official PyTorch website.

2. Step 2. Organise the dataset into class folders

   ImageFolder expects one subfolder per class, with the folder name as the label.

   data/
     tomato_early_blight/
       img001.jpg
       img002.jpg
     tomato_healthy/
       img001.jpg
     tomato_leaf_mold/
       img001.jpg

   For a plant disease project, PlantVillage is the usual starting point: TensorFlow Datasets describes it as 54,303 leaf images in 38 classes by species and disease, and notes the original source site no longer hosts it, so use a mirror such as the one listed there or a Kaggle copy. Start with 3 to 5 classes from one crop. Fewer, well-chosen classes give a cleaner project than all 38 at once, and train much faster.

   Check the class balance before training. If one class has ten times the images of another, accuracy will flatter the model, and per-class recall becomes the number to watch.

3. Step 3. Load the data with transforms and a train/validation split

   Training images get random augmentation; validation images get only the standard resize, crop and normalisation that the pretrained weights expect.

   # data.py
   import torch
   from torch.utils.data import DataLoader, Subset
   from torchvision import datasets, transforms
   from torchvision.models import ResNet18_Weights
   
   MEAN, STD = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]  # ImageNet statistics
   train_tf = transforms.Compose([
       transforms.RandomResizedCrop(224),
       transforms.RandomHorizontalFlip(),
       transforms.ColorJitter(0.2, 0.2, 0.2),
       transforms.ToTensor(),
       transforms.Normalize(MEAN, STD),
   ])
   val_tf = ResNet18_Weights.DEFAULT.transforms()  # resize, center crop, normalise
   
   def get_loaders(root="data", val_frac=0.2, batch_size=32):
       train_full = datasets.ImageFolder(root, transform=train_tf)
       val_full = datasets.ImageFolder(root, transform=val_tf)
       g = torch.Generator().manual_seed(42)
       idx = torch.randperm(len(train_full), generator=g).tolist()
       n_val = int(len(idx) * val_frac)
       train_ds, val_ds = Subset(train_full, idx[n_val:]), Subset(val_full, idx[:n_val])
       return (DataLoader(train_ds, batch_size=batch_size, shuffle=True, num_workers=2),
               DataLoader(val_ds, batch_size=batch_size, num_workers=2),
               train_full.classes)

   Two ImageFolder objects over the same folder, split with the same shuffled indices, let each split use its own transforms. The fixed seed means every script sees the same validation images.

4. Step 4. Load a pretrained ResNet18 and replace the head

   torchvision downloads the ImageNet weights on first use. We freeze all layers, then swap the final fully connected layer for one with an output per class.

   # model.py
   import torch.nn as nn
   from torchvision.models import resnet18, ResNet18_Weights
   
   def build_model(num_classes, freeze=True):
       model = resnet18(weights=ResNet18_Weights.DEFAULT)
       if freeze:
           for p in model.parameters():
               p.requires_grad = False
       model.fc = nn.Linear(model.fc.in_features, num_classes)  # new layer, trainable
       return model

   The new layer is created after freezing, so it is the only part with requires_grad set to True. ResNet18 has about 11.7 million parameters (listed in the torchvision docs); in this phase you train only a few thousand of them.

5. Step 5. Train the new layer and save the best model

   One function handles both training and validation, so the two loops cannot drift apart.

   # train.py
   import torch
   import torch.nn as nn
   from data import get_loaders
   from model import build_model
   
   def run_epoch(model, loader, loss_fn, device, optimizer=None):
       training = optimizer is not None
       model.train(training)
       total, correct, loss_sum = 0, 0, 0.0
       with torch.set_grad_enabled(training):
           for x, y in loader:
               x, y = x.to(device), y.to(device)
               out = model(x)
               loss = loss_fn(out, y)
               if training:
                   optimizer.zero_grad()
                   loss.backward()
                   optimizer.step()
               loss_sum += loss.item() * len(y)
               correct += (out.argmax(1) == y).sum().item()
               total += len(y)
       return loss_sum / total, correct / total
   
   if __name__ == "__main__":
       device = "cuda" if torch.cuda.is_available() else "cpu"
       train_dl, val_dl, classes = get_loaders()
       model = build_model(len(classes)).to(device)
       optimizer = torch.optim.Adam(model.fc.parameters(), lr=1e-3)
       loss_fn, best = nn.CrossEntropyLoss(), 0.0
       for epoch in range(5):
           _, tr_acc = run_epoch(model, train_dl, loss_fn, device, optimizer)
           va_loss, va_acc = run_epoch(model, val_dl, loss_fn, device)
           print(f"epoch {epoch + 1}: train {tr_acc:.3f}  val {va_acc:.3f}  val loss {va_loss:.3f}")
           if va_acc > best:
               best = va_acc
               torch.save({"state": model.state_dict(), "classes": classes}, "best_model.pt")

   The if __name__ == "__main__" guard is needed on Windows when DataLoader uses worker processes. Log the printed numbers for your report, and plot them if you have time: the gap between train and val accuracy tells the overfitting story.

6. Step 6. Fine-tune the last block for extra accuracy

   Once the new layer has learned, unfreeze the last residual block (layer4) and train both with a smaller learning rate for that block. In train.py, after the first loop finishes, add:

   for p in model.layer4.parameters():
       p.requires_grad = True
   
   optimizer = torch.optim.Adam([
       {"params": model.layer4.parameters(), "lr": 1e-4},
       {"params": model.fc.parameters(), "lr": 1e-3},
   ])
   for epoch in range(5):
       _, tr_acc = run_epoch(model, train_dl, loss_fn, device, optimizer)
       va_loss, va_acc = run_epoch(model, val_dl, loss_fn, device)
       print(f"fine-tune {epoch + 1}: train {tr_acc:.3f}  val {va_acc:.3f}")
       if va_acc > best:
           best = va_acc
           torch.save({"state": model.state_dict(), "classes": classes}, "best_model.pt")

   The smaller learning rate stops large updates from wrecking the pretrained features. Report validation accuracy after phase 1 and after phase 2 separately: that comparison is a ready-made results table.

7. Step 7. Evaluate with a classification report and confusion matrix

   Accuracy is one number. The confusion matrix shows which classes the model confuses, which is where the real insight sits.

   # evaluate.py
   import torch
   import matplotlib.pyplot as plt
   from sklearn.metrics import classification_report, ConfusionMatrixDisplay
   from data import get_loaders
   from model import build_model
   
   if __name__ == "__main__":
       device = "cuda" if torch.cuda.is_available() else "cpu"
       _, val_dl, classes = get_loaders()
       ckpt = torch.load("best_model.pt", map_location=device)
       model = build_model(len(classes)).to(device)
       model.load_state_dict(ckpt["state"])
       model.eval()
   
       preds, labels = [], []
       with torch.no_grad():
           for x, y in val_dl:
               preds += model(x.to(device)).argmax(1).cpu().tolist()
               labels += y.tolist()
   
       print(classification_report(labels, preds, target_names=classes))
       ConfusionMatrixDisplay.from_predictions(labels, preds, display_labels=classes,
                                               xticks_rotation=45)
       plt.tight_layout()
       plt.savefig("confusion_matrix.png", dpi=150)

   Put confusion_matrix.png in your report and explain the largest off-diagonal cell. "Early blight and late blight get confused because both show brown lesions" is the kind of sentence that earns marks.

8. Step 8. Predict on a new photo

   A prediction script turns the model into something you can demo with a fresh photo taken on your phone.

   # predict.py
   import sys
   import torch
   from PIL import Image
   from torchvision.models import ResNet18_Weights
   from model import build_model
   
   ckpt = torch.load("best_model.pt", map_location="cpu")
   model = build_model(len(ckpt["classes"]))
   model.load_state_dict(ckpt["state"])
   model.eval()
   tf = ResNet18_Weights.DEFAULT.transforms()
   
   img = Image.open(sys.argv[1]).convert("RGB")
   with torch.no_grad():
       probs = torch.softmax(model(tf(img).unsqueeze(0)), dim=1)[0]
   top = probs.topk(min(3, len(probs)))
   for p, i in zip(top.values.tolist(), top.indices.tolist()):
       print(f"{ckpt['classes'][i]}: {p:.1%}")

   Run python predict.py leaf.jpg. Test it on photos taken outside the dataset. The drop in confidence you usually see on real-world photos is worth a paragraph in your limitations section.

Frequently asked questions

Keep learning

References

• PyTorch tutorial: Transfer learning for computer vision
• torchvision: models and pretrained weights
• torchvision: ImageFolder documentation
• TensorFlow Datasets: PlantVillage catalog entry
• Google Colab FAQ