Transfer Learning Flower Classifier

Fine-tune a pretrained ResNet18 to classify 102 flower types with PyTorch.

Introduction

30 Second Summary

Apps like Google Photos and plant identification tools can recognize thousands of objects from a single photo. They achieve this by building on visual patterns already learned from millions of images, not by starting from scratch.

In this project, you will build a flower image classifier that reuses a pretrained neural network to identify 102 types of flowers, achieving strong accuracy on your CPU alone.

What You'll Build

A flower image classifier powered by a pretrained ResNet18 model, fine-tuned on the Oxford 102 Flowers dataset using PyTorch.

By the end of this project, you'll have:

  • A complete PyTorch training pipeline with data loading, image transforms, a training loop, and validation.
  • A trained transfer learning model that classifies 102 flower categories with strong accuracy, running entirely on CPU.
  • A working inference script that predicts flower types from new images and displays top-5 predictions with confidence scores.
  • Secret Mission: Implement gradual layer unfreezing to push accuracy beyond the frozen-backbone approach.

Are there any prerequisites?

This project assumes intermediate Python knowledge and a basic understanding of machine learning concepts like training loops, loss functions, and optimizers. You do not need prior experience with PyTorch or computer vision.

Load and Explore the Dataset

Your environment is set up and ready to go. Now it's time to bring in the data your model will learn from.

A deep learning model is only as good as the data it trains on. In this step, you'll download the Oxford Flowers 102 dataset and build a transform pipeline that prepares each image for your neural network.

In this step, get ready to:
  • Define image transforms for training and validation.
  • Download the Oxford Flowers 102 dataset.
  • Create DataLoaders that feed batches of images to your model.
Define the image transforms

Neural networks expect input images to have a consistent size and pixel range. Transforms handle this preprocessing.

They also add random variations during training to help the model generalize better. This technique is called data augmentation.

  • Open train.py in VS Code and add the following imports at the top of the file:
import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim.lr_scheduler import StepLR
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import Flowers102
from torchvision.models import resnet18, ResNet18_Weights

What are these imports?

  • torch is the core PyTorch library for tensor operations and neural network building.
  • DataLoader and transforms handle batching images and preprocessing them.
  • Flowers102 is the dataset class that downloads and loads the flower images.
  • The remaining imports (nn, optim, StepLR, resnet18, ResNet18_Weights) will be used in the next step when you build the model.
  • Save train.py.
  • You won't see any output yet. These imports prepare the building blocks for the functions below.

Seeing a ModuleNotFoundError?

  • Make sure your virtual environment is activated. You should see (venv) at the start of your terminal prompt.
  • Check that you installed torchvision in Step 1. If not, reinstall the packages.

Help me fix a ModuleNotFoundError in my PyTorch project

  • Below the imports, add the get_transforms() function that defines how images are preprocessed:
def get_transforms():
    train_transform = transforms.Compose([
        transforms.Resize(256),
        transforms.RandomCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225]
        ),
    ])

    val_transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225]
        ),
    ])

    return train_transform, val_transform

What does this code do?

This function creates two transform pipelines. The training pipeline randomizes how images are cropped and flipped, so the model sees slightly different versions each epoch. The validation pipeline uses a fixed center crop so evaluation results stay consistent.

  • Resize(256) scales the image so its shorter side is 256 pixels.
  • RandomCrop(224) takes a random 224x224 patch during training. CenterCrop(224) takes the center patch during validation.
  • Normalize applies the ImageNet mean and standard deviation values. These numbers are required because the pretrained ResNet18 was trained with them.
  • Save train.py.
  • Run python train.py in the terminal. You should see no output and no errors.

Seeing a SyntaxError?

  • Check that your parentheses and brackets are balanced. Each opening [ and ( needs a matching close.
  • Make sure each transform line inside the list ends with a comma.

Help me fix a SyntaxError in my transforms code

Download the dataset and create DataLoaders

The Oxford Flowers 102 dataset contains images across 102 flower categories commonly found in the United Kingdom. Each category has between 40 and 258 images with large variations in scale, pose, and lighting.

A DataLoader is a PyTorch utility that groups images into batches and shuffles them during training. This is how your model receives data efficiently.

  • Below the get_transforms() function in train.py, add the load_data() function:
def load_data(train_transform, val_transform):
    train_dataset = Flowers102(
        root="data",
        split="train",
        transform=train_transform,
        download=True,
    )

    val_dataset = Flowers102(
        root="data",
        split="val",
        transform=val_transform,
        download=True,
    )

    train_loader = DataLoader(
        train_dataset, batch_size=32, shuffle=True, num_workers=0
    )
    val_loader = DataLoader(
        val_dataset, batch_size=32, shuffle=False, num_workers=0
    )

    return train_loader, val_loader

What does this code do?

  • Flowers102(..., download=True) downloads the image archive the first time you run it. Subsequent runs skip the download.
  • split="train" and split="val" load separate subsets so you can train on one and evaluate on the other.
  • batch_size=32 means the DataLoader feeds 32 images at a time to the model. This balances memory usage with training speed.
  • num_workers=0 disables multiprocessing for data loading. On Windows, multiprocessing in DataLoaders causes errors without this setting.

Why num_workers=0 on Windows?

PyTorch's DataLoader uses Python multiprocessing to load data in parallel. On Windows, this triggers spawn-based process creation that conflicts with how Python modules are imported.

Setting num_workers=0 loads data in the main process, avoiding the issue entirely.

  • Save train.py.
  • Verify that both functions work and trigger the dataset download by running this command:
python -c "from train import get_transforms, load_data; t, v = get_transforms(); tl, vl = load_data(t, v); print('Train batches:', len(tl), 'Val batches:', len(vl))"
  • The first run downloads the dataset, so you will see download progress messages. This may take a minute depending on your internet speed.
  • Once the download finishes, you should see output showing the number of train and validation batches (around 32 train batches and 32 val batches).

Seeing an error about scipy?

The Flowers102 dataset requires scipy to load its label files. If you see an import error mentioning scipy, install it by running:

pip install scipy

If you see a network timeout or connection error, check your internet connection and try the command again.

Help me fix a dataset download error in my PyTorch project

✔️ Awesome, I've got everything!

Great. Double check you've saved train.py and that your data/ folder now contains the downloaded flower images.

ⓧ I'd like to double check the full code

Here is the complete train.py so far. Make sure yours matches:

import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim.lr_scheduler import StepLR
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import Flowers102
from torchvision.models import resnet18, ResNet18_Weights


def get_transforms():
    train_transform = transforms.Compose([
        transforms.Resize(256),
        transforms.RandomCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225]
        ),
    ])

    val_transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225]
        ),
    ])

    return train_transform, val_transform


def load_data(train_transform, val_transform):
    train_dataset = Flowers102(
        root="data",
        split="train",
        transform=train_transform,
        download=True,
    )

    val_dataset = Flowers102(
        root="data",
        split="val",
        transform=val_transform,
        download=True,
    )

    train_loader = DataLoader(
        train_dataset, batch_size=32, shuffle=True, num_workers=0
    )
    val_loader = DataLoader(
        val_dataset, batch_size=32, shuffle=False, num_workers=0
    )

    return train_loader, val_loader

Your data pipeline is ready. Next up, you'll build the transfer learning model that uses a pretrained ResNet18 backbone to classify these flowers.

Build the Transfer Learning Model

Your dataset is loaded and ready to feed into a model. Now you need to build the model itself.

Training a deep neural network from scratch on only 1,020 images would produce terrible results. Instead, you will use transfer learning to take a model that already understands visual features (edges, textures, shapes) and repurpose it for flower classification.

In this step, get ready to:
  • Load a pretrained ResNet18 model and freeze its backbone layers.
  • Write a training loop function that handles forward passes, loss computation, and backpropagation.
  • Write a validation function that evaluates model performance without updating weights.
Load pretrained ResNet18 and freeze its backbone

ResNet18 is a convolutional neural network pretrained on ImageNet, a dataset of over 1 million images across 1,000 categories. Its convolutional layers already know how to detect visual features like edges, colors, and textures.

By freezing these layers (preventing their weights from updating), you keep all that learned knowledge intact. Then you replace just the final classification layer so the model outputs 102 flower categories instead of 1,000 ImageNet categories.

  • In train.py, find the end of the load_data function.
  • Add the following create_model function below it:
def create_model(num_classes=102):
    # Load ResNet18 with pretrained ImageNet weights
    model = resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)

    # Freeze all layers so their weights won't update during training
    for param in model.parameters():
        param.requires_grad = False

    # Replace the final classification layer (512 inputs -> 102 flower classes)
    num_ftrs = model.fc.in_features
    model.fc = nn.Linear(num_ftrs, num_classes)

    return model

What does this code do?

This function builds your transfer learning model in three steps:

  • resnet18(weights=ResNet18_Weights.IMAGENET1K_V1) downloads the pretrained model with weights trained on ImageNet.
  • param.requires_grad = False freezes every layer so gradients won't flow through them during training. This preserves the learned visual features.
  • model.fc = nn.Linear(num_ftrs, num_classes) replaces the original 1000-class output layer with a fresh 102-class layer. Because this new layer was not frozen, it is the only part that will learn.

Why freeze instead of training everything?

ResNet18 has 11.7 million parameters. Training all of them on only 1,020 flower images would lead to severe overfitting. Freezing the backbone means you only train the 52,326 parameters in the new classification layer.

This makes training fast (even on CPU) and prevents the model from forgetting the visual features it already learned from ImageNet.

  • Save train.py.

Seeing a DeprecationWarning about pretrained?

If you used pretrained=True instead of weights=ResNet18_Weights.IMAGENET1K_V1, you will see a deprecation warning. Make sure you are using the weights parameter as shown above.

help me fix the pretrained parameter deprecation warning in my ResNet18 model

Write the training loop function

Your model needs a function that handles one full pass through the training data. Each pass feeds batches of images forward, computes how wrong the predictions are, and updates the classification layer to be less wrong next time.

  • In train.py, add the following train_one_epoch function below create_model:
def train_one_epoch(model, train_loader, criterion, optimizer, device):
    # Set model to training mode (enables dropout, batch norm uses batch stats)
    model.train()
    running_loss = 0.0
    correct = 0
    total = 0

    for images, labels in train_loader:
        # Move data to the correct device
        images, labels = images.to(device), labels.to(device)

        # Zero out old gradients before computing new ones
        optimizer.zero_grad()
        # Forward pass: get predictions
        outputs = model(images)
        # Compute how wrong the predictions are
        loss = criterion(outputs, labels)
        # Backward pass: compute gradients
        loss.backward()
        # Update the trainable weights
        optimizer.step()

        # Track metrics for reporting
        running_loss += loss.item() * images.size(0)
        _, predicted = torch.max(outputs, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

    epoch_loss = running_loss / total
    epoch_acc = correct / total
    return epoch_loss, epoch_acc

What does this code do?

This function runs one complete epoch of training. Here is the core loop:

  • model.train() switches the model to training mode, which activates batch normalization's batch-statistics behavior.
  • optimizer.zero_grad() clears leftover gradients from the previous batch so they don't accumulate.
  • loss.backward() computes how much each trainable weight contributed to the error.
  • optimizer.step() nudges the weights in the direction that reduces the loss.
  • The function returns average loss and accuracy so you can track progress across epochs.
  • Save train.py.

Unsure about the training loop structure?

The order matters: zero gradients, forward pass, compute loss, backward pass, optimizer step. If you call optimizer.step() before loss.backward(), the weights update using stale gradients from a previous batch.

explain the correct order of operations in a PyTorch training loop and why it matters

Write the validation function

Training accuracy alone can be misleading because the model might just memorize the training images. You need a validation function that tests performance on images the model has never trained on.

  • In train.py, add the following validate function below train_one_epoch:
def validate(model, val_loader, criterion, device):
    # Set model to evaluation mode (disables dropout, batch norm uses running stats)
    model.eval()
    running_loss = 0.0
    correct = 0
    total = 0

    # Disable gradient computation to save memory during validation
    with torch.no_grad():
        for images, labels in val_loader:
            images, labels = images.to(device), labels.to(device)

            outputs = model(images)
            loss = criterion(outputs, labels)

            running_loss += loss.item() * images.size(0)
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    epoch_loss = running_loss / total
    epoch_acc = correct / total
    return epoch_loss, epoch_acc

What does this code do?

This function evaluates the model without changing any weights:

  • model.eval() switches to evaluation mode. Batch normalization layers use their stored running statistics instead of batch statistics, giving consistent results.
  • torch.no_grad() disables gradient tracking entirely. This saves memory because PyTorch does not need to store intermediate values for backpropagation.
  • The function returns validation loss and accuracy so you can compare against training metrics and detect overfitting.
  • Save train.py.

Confused about model.train() vs model.eval()?

These two calls change how certain layers behave. In training mode, batch normalization uses the current batch's statistics. In evaluation mode, it uses stable running averages collected during training. Forgetting to call model.eval() before validation gives inconsistent accuracy numbers.

explain the difference between model.train() and model.eval() for my validation function

✔️ Awesome, I've got everything!

Great. Double check you have saved train.py with all three new functions.

ⓧ I'd like to double check the full code

Here is the complete train.py so far:

import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim.lr_scheduler import StepLR
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import Flowers102
from torchvision.models import resnet18, ResNet18_Weights


def get_transforms():
    train_transform = transforms.Compose([
        transforms.Resize(256),
        transforms.RandomCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225]
        ),
    ])

    val_transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225]
        ),
    ])

    return train_transform, val_transform


def load_data(train_transform, val_transform):
    train_dataset = Flowers102(
        root="data",
        split="train",
        transform=train_transform,
        download=True,
    )

    val_dataset = Flowers102(
        root="data",
        split="val",
        transform=val_transform,
        download=True,
    )

    train_loader = DataLoader(
        train_dataset, batch_size=32, shuffle=True, num_workers=0
    )
    val_loader = DataLoader(
        val_dataset, batch_size=32, shuffle=False, num_workers=0
    )

    return train_loader, val_loader


def create_model(num_classes=102):
    model = resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)

    for param in model.parameters():
        param.requires_grad = False

    num_ftrs = model.fc.in_features
    model.fc = nn.Linear(num_ftrs, num_classes)

    return model


def train_one_epoch(model, train_loader, criterion, optimizer, device):
    model.train()
    running_loss = 0.0
    correct = 0
    total = 0

    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)

        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item() * images.size(0)
        _, predicted = torch.max(outputs, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

    epoch_loss = running_loss / total
    epoch_acc = correct / total
    return epoch_loss, epoch_acc


def validate(model, val_loader, criterion, device):
    model.eval()
    running_loss = 0.0
    correct = 0
    total = 0

    with torch.no_grad():
        for images, labels in val_loader:
            images, labels = images.to(device), labels.to(device)

            outputs = model(images)
            loss = criterion(outputs, labels)

            running_loss += loss.item() * images.size(0)
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    epoch_loss = running_loss / total
    epoch_acc = correct / total
    return epoch_loss, epoch_acc

Your model architecture is ready. Next up, you will wire these functions together into a training loop that runs for multiple epochs and saves the best-performing model.

Train and Validate

You have all the building blocks in place. Your train.py file knows how to load data, build a model, train for one epoch, and validate.

Now you need to orchestrate these pieces into a complete training pipeline. The main() function will wire everything together, run the model through 10 epochs, and save the best-performing checkpoint to disk.

In this step, get ready to:
  • Write the training configuration and epoch loop.
  • Add validation tracking and model checkpointing.
  • Run the full training for 10 epochs and save the best model.
Set up the main function with training configuration

The main() function is where all your helper functions come together. It sets the device, loads data, creates the model, and configures the loss function, optimizer, and learning rate scheduler.

  • In train.py, scroll to the very bottom of the file (below the validate() function).
  • Add the following main() function:
def main():
    # Set device to CPU (no GPU needed for this project)
    device = torch.device("cpu")
    print(f"Using device: {device}")

    # Load transforms and data using our helper functions
    train_transform, val_transform = get_transforms()
    train_loader, val_loader = load_data(train_transform, val_transform)
    print(f"Training samples: {len(train_loader.dataset)}")
    print(f"Validation samples: {len(val_loader.dataset)}")

    # Create the transfer learning model and move to device
    model = create_model(num_classes=102)
    model = model.to(device)

    # Set up loss function, optimizer (only training fc layer), and scheduler
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.fc.parameters(), lr=0.001)
    scheduler = StepLR(optimizer, step_size=5, gamma=0.1)

    # Training configuration
    num_epochs = 10
    best_val_acc = 0.0

What does this code do?

This setup section orchestrates the pieces you built in previous steps.

  • torch.device("cpu") tells PyTorch to run all computations on the CPU.
  • optim.Adam(model.fc.parameters(), lr=0.001) creates an optimizer that only updates the new classification head. The frozen backbone parameters are excluded.
  • StepLR(optimizer, step_size=5, gamma=0.1) reduces the learning rate by 10x every 5 epochs. This helps the model converge more precisely in later epochs.
  • best_val_acc = 0.0 tracks the highest validation accuracy seen so far, so you only save the best-performing model.

Why Adam instead of SGD?

Adam converges faster with default hyperparameters, which matters when your training set is small (only ~10 images per class in Flowers102 train split). SGD requires more careful tuning of momentum and learning rate to get the same convergence speed.

Add the training and validation loop

The epoch loop calls your train_one_epoch() and validate() functions each iteration, prints metrics, and saves the model whenever validation accuracy improves.

  • Directly below the best_val_acc = 0.0 line you just added, paste the following code (still inside main()):
    for epoch in range(num_epochs):
        train_loss, train_acc = train_one_epoch(
            model, train_loader, criterion, optimizer, device
        )
        val_loss, val_acc = validate(model, val_loader, criterion, device)
        scheduler.step()

        print(
            f"Epoch [{epoch+1}/{num_epochs}] "
            f"Train Loss: {train_loss:.4f} Acc: {train_acc:.4f} | "
            f"Val Loss: {val_loss:.4f} Acc: {val_acc:.4f}"
        )

        if val_acc > best_val_acc:
            best_val_acc = val_acc
            torch.save(model.state_dict(), "best_model.pth")
            print(f"  -> Saved best model (Val Acc: {val_acc:.4f})")

    print(f"\nTraining complete. Best validation accuracy: {best_val_acc:.4f}")


if __name__ == "__main__":
    main()

What does this code do?

  • The for epoch in range(num_epochs) loop runs the full training cycle 10 times. Each epoch processes every batch in the training set.
  • After each training pass, validate() evaluates the model on held-out data without updating weights. This tells you whether the model is actually learning to generalize.
  • scheduler.step() decreases the learning rate on schedule.
  • torch.save(model.state_dict(), "best_model.pth") saves only the model weights (not the full model object). This is the recommended approach for model persistence in PyTorch.
  • The model is only saved when validation accuracy improves. This technique is called checkpointing.
  • if __name__ == "__main__": ensures main() only runs when you execute train.py directly (not when importing from it).
  • Save train.py.

✔️ Awesome, I've got everything!

Great. Double check you have saved train.py before moving on.

ⓧ I'd like to double check the full code

Here is the complete train.py file. Compare yours against it:

import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim.lr_scheduler import StepLR
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import Flowers102
from torchvision.models import resnet18, ResNet18_Weights


def get_transforms():
    train_transform = transforms.Compose([
        transforms.Resize(256),
        transforms.RandomCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225]
        ),
    ])

    val_transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225]
        ),
    ])

    return train_transform, val_transform


def load_data(train_transform, val_transform):
    train_dataset = Flowers102(
        root="data",
        split="train",
        transform=train_transform,
        download=True,
    )

    val_dataset = Flowers102(
        root="data",
        split="val",
        transform=val_transform,
        download=True,
    )

    train_loader = DataLoader(
        train_dataset, batch_size=32, shuffle=True, num_workers=0
    )
    val_loader = DataLoader(
        val_dataset, batch_size=32, shuffle=False, num_workers=0
    )

    return train_loader, val_loader


def create_model(num_classes=102):
    model = resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)

    for param in model.parameters():
        param.requires_grad = False

    num_ftrs = model.fc.in_features
    model.fc = nn.Linear(num_ftrs, num_classes)

    return model


def train_one_epoch(model, train_loader, criterion, optimizer, device):
    model.train()
    running_loss = 0.0
    correct = 0
    total = 0

    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)

        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item() * images.size(0)
        _, predicted = torch.max(outputs, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

    epoch_loss = running_loss / total
    epoch_acc = correct / total
    return epoch_loss, epoch_acc


def validate(model, val_loader, criterion, device):
    model.eval()
    running_loss = 0.0
    correct = 0
    total = 0

    with torch.no_grad():
        for images, labels in val_loader:
            images, labels = images.to(device), labels.to(device)

            outputs = model(images)
            loss = criterion(outputs, labels)

            running_loss += loss.item() * images.size(0)
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    epoch_loss = running_loss / total
    epoch_acc = correct / total
    return epoch_loss, epoch_acc


def main():
    device = torch.device("cpu")
    print(f"Using device: {device}")

    train_transform, val_transform = get_transforms()
    train_loader, val_loader = load_data(train_transform, val_transform)
    print(f"Training samples: {len(train_loader.dataset)}")
    print(f"Validation samples: {len(val_loader.dataset)}")

    model = create_model(num_classes=102)
    model = model.to(device)

    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.fc.parameters(), lr=0.001)
    scheduler = StepLR(optimizer, step_size=5, gamma=0.1)

    num_epochs = 10
    best_val_acc = 0.0

    for epoch in range(num_epochs):
        train_loss, train_acc = train_one_epoch(
            model, train_loader, criterion, optimizer, device
        )
        val_loss, val_acc = validate(model, val_loader, criterion, device)
        scheduler.step()

        print(
            f"Epoch [{epoch+1}/{num_epochs}] "
            f"Train Loss: {train_loss:.4f} Acc: {train_acc:.4f} | "
            f"Val Loss: {val_loss:.4f} Acc: {val_acc:.4f}"
        )

        if val_acc > best_val_acc:
            best_val_acc = val_acc
            torch.save(model.state_dict(), "best_model.pth")
            print(f"  -> Saved best model (Val Acc: {val_acc:.4f})")

    print(f"\nTraining complete. Best validation accuracy: {best_val_acc:.4f}")


if __name__ == "__main__":
    main()
Run training for 10 epochs

Time to run the training pipeline. This will take a few minutes on CPU since the model processes all training batches 10 times.

  • Run the training script in your terminal:
python train.py

What should I see?

You will see output like this for each epoch:

Using device: cpu
Training samples: 1020
Validation samples: 1020
Epoch [1/10] Train Loss: 4.2538 Acc: 0.0814 | Val Loss: 3.8021 Acc: 0.1245
  -> Saved best model (Val Acc: 0.1245)
Epoch [2/10] Train Loss: 3.5102 Acc: 0.1892 | Val Loss: 3.2156 Acc: 0.2304
  -> Saved best model (Val Acc: 0.2304)
...

Your exact numbers will differ slightly due to random initialization. Training accuracy will climb each epoch, and the script saves the model whenever validation accuracy hits a new high.

Training takes a few minutes

On CPU, expect each epoch to take 30-60 seconds depending on your machine. The full 10-epoch run typically finishes in 5-10 minutes. You can watch accuracy improve in real time as each epoch prints.

  • Wait for all 10 epochs to complete. You should see a final line showing the best validation accuracy.
  • Confirm that best_model.pth was saved by running:
ls best_model.pth

You should see best_model.pth listed in the output. This file contains the trained weights you will use for inference in the next step.

Training seems stuck or errors out?

  • If you see an indentation error, make sure the entire for epoch loop is indented inside main(). Every line inside main() needs 4 spaces of indentation.
  • If nothing happens when you run the script, check that if __name__ == "__main__": is at the very bottom of the file with no indentation.
  • If you see a NameError, make sure you have not accidentally deleted any of the import statements at the top of train.py.

Still stuck? Help me debug my train.py script. I'm trying to run 10 epochs of training but it's not working.

Your model is trained and saved. Next up, you will build an inference script that loads this model and classifies new flower images with confidence scores.

Run Inference on New Images

Your model is trained and saved. The real test of any classifier is whether it works on images it has never seen before.

In this step, you will build an inference script that loads your saved model, preprocesses a flower image, and outputs the top-5 predicted flower names with confidence percentages.

In this step, get ready to:
  • Write a prediction function that loads the saved model and classifies a single image.
  • Extract top-5 predictions with confidence scores.
  • Display the image alongside predicted class names using matplotlib.
Set up imports and the flower class names

The prediction script needs to map numeric class indices back to human-readable flower names. You will define these at the top of the file alongside the imports.

  • Open predict.py in VS Code and add the imports and flower names list by pasting the following code:
import sys

import matplotlib.pyplot as plt
import torch
import torch.nn as nn
from torchvision import transforms
from torchvision.models import resnet18
from PIL import Image


# Human-readable names for each of the 102 flower classes
FLOWER_NAMES = [
    "pink primrose", "hard-leaved pocket orchid", "canterbury bells",
    "sweet pea", "english marigold", "tiger lily", "moon orchid",
    "bird of paradise", "monkshood", "globe thistle", "snapdragon",
    "colt's foot", "king protea", "spear thistle", "yellow iris",
    "globe-flower", "purple coneflower", "peruvian lily", "balloon flower",
    "giant white arum lily", "fire lily", "pincushion flower",
    "fritillary", "red ginger", "grape hyacinth", "corn poppy",
    "prince of wales feathers", "stemless gentian", "artichoke",
    "sweet william", "carnation", "garden phlox", "love in the mist",
    "mexican aster", "alpine sea holly", "ruby-lipped cattleya",
    "cape flower", "great masterwort", "siam tulip", "lenten rose",
    "barbeton daisy", "daffodil", "sword lily", "poinsettia",
    "bolero deep blue", "wallflower", "marigold", "buttercup",
    "oxeye daisy", "common dandelion", "petunia", "wild pansy",
    "primula", "sunflower", "pelargonium", "bishop of llandaff",
    "gaura", "geranium", "orange dahlia", "pink-yellow dahlia",
    "cautleya spicata", "japanese anemone", "black-eyed susan",
    "silverbush", "californian poppy", "osteospermum", "spring crocus",
    "bearded iris", "windflower", "tree poppy", "gazania",
    "azalea", "water lily", "rose", "thorn apple", "morning glory",
    "passion flower", "lotus", "toad lily", "anthurium",
    "frangipani", "clematis", "hibiscus", "columbine", "desert-rose",
    "tree mallow", "magnolia", "cyclamen", "watercress",
    "canna lily", "hippeastrum", "bee balm", "ball moss",
    "foxglove", "bougainvillea", "camellia", "mallow",
    "mexican petunia", "bromelia", "blanket flower",
    "trumpet creeper", "blackberry lily",
]

What does this code do?

The imports bring in PyTorch for model loading, torchvision for image transforms, PIL for reading image files, and matplotlib for visualizing results.

The FLOWER_NAMES list maps each numeric class index (0-101) to its flower name. When the model predicts class 53, you look up FLOWER_NAMES[53] to get the human-readable name.

Write the model loading function

At inference time, you rebuild the same model architecture but load the trained weights from disk instead of the pretrained ImageNet weights.

  • Below the FLOWER_NAMES list, add the load_model function:
def load_model(model_path, num_classes=102):
    # Rebuild the same architecture without pretrained weights
    model = resnet18(weights=None)
    num_ftrs = model.fc.in_features
    model.fc = nn.Linear(num_ftrs, num_classes)
    # Load the trained weights from disk
    model.load_state_dict(torch.load(model_path, map_location="cpu"))
    model.eval()
    return model

What does this code do?

The function creates a fresh ResNet18 with weights=None (no pretrained weights), replaces the final layer to match 102 classes, then loads your trained weights from best_model.pth.

Calling model.eval() switches the model to evaluation mode. This disables dropout and batch normalization training behavior so predictions are deterministic.

Write the prediction function with top-5 confidence scores

The prediction function preprocesses a single image using the same transforms the model was validated on, runs a forward pass, and extracts the top-5 most confident predictions.

  • Below load_model, add the predict_image function:
def predict_image(model, image_path):
    # Same transforms used during validation
    transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225]
        ),
    ])

    # Load and preprocess the image
    image = Image.open(image_path).convert("RGB")
    input_tensor = transform(image).unsqueeze(0)

    # Run forward pass and extract top-5 predictions
    with torch.no_grad():
        outputs = model(input_tensor)
        probabilities = torch.softmax(outputs, dim=1)
        top5_prob, top5_idx = torch.topk(probabilities, 5)

    # Map indices to flower names
    results = []
    for i in range(5):
        idx = top5_idx[0][i].item()
        prob = top5_prob[0][i].item()
        name = FLOWER_NAMES[idx] if idx < len(FLOWER_NAMES) else f"Class {idx}"
        results.append((name, prob))

    return image, results

What does this code do?

The validation transforms (Resize to 256, CenterCrop to 224, normalize with ImageNet statistics) must match what the model saw during training. Using different transforms would produce garbage predictions.

The .unsqueeze(0) call adds a batch dimension because PyTorch models expect input shaped as (batch, channels, height, width).

Inside torch.no_grad(), the model runs a forward pass without tracking gradients (saving memory). torch.softmax converts raw logits into probabilities that sum to 1. torch.topk returns the 5 highest probabilities and their indices.

  • Save predict.py.

You will not see a visible change yet. The next function uses these results to display a chart.

Display predictions with matplotlib

A bar chart showing the model's confidence for each top prediction makes it easy to interpret results at a glance.

  • Below predict_image, add the display_prediction function:
def display_prediction(image, results, image_path):
    # Create a side-by-side figure: image on the left, bar chart on the right
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))

    ax1.imshow(image)
    ax1.set_title(f"Input: {image_path}")
    ax1.axis("off")

    # Extract names and percentages for the bar chart
    names = [r[0] for r in results]
    probs = [r[1] * 100 for r in results]

    ax2.barh(range(5), probs, color="steelblue")
    ax2.set_yticks(range(5))
    ax2.set_yticklabels(names)
    ax2.set_xlabel("Confidence (%)")
    ax2.set_title("Top 5 Predictions")
    ax2.set_xlim(0, 100)

    plt.tight_layout()
    plt.savefig("prediction_result.png", dpi=100, bbox_inches="tight")
    plt.show()
    print("Prediction saved to prediction_result.png")

What does this code do?

The function creates a matplotlib figure with two panels. The left panel shows the original flower image. The right panel displays a horizontal bar chart of the top-5 predicted flower names and their confidence percentages.

The chart is saved to prediction_result.png so you have a permanent record of the prediction even after closing the window.

  • Now add the main function that ties everything together:
def main():
    # Check that the user provided an image path as a command-line argument
    if len(sys.argv) < 2:
        print("Usage: python predict.py <image_path>")
        print("Example: python predict.py data/flowers-102/jpg/image_00001.jpg")
        sys.exit(1)

    image_path = sys.argv[1]
    model = load_model("best_model.pth")
    image, results = predict_image(model, image_path)

    # Print predictions to the terminal
    print("\nPredictions:")
    for name, prob in results:
        print(f"  {name}: {prob*100:.1f}%")

    display_prediction(image, results, image_path)


if __name__ == "__main__":
    main()

What does this code do?

The main function reads the image file path from the command line (sys.argv[1]), loads the trained model, runs the prediction, prints text results to the terminal, and opens the visualization window.

The if __name__ == "__main__" guard ensures main() only runs when you execute the script directly, not when importing it.

  • Save predict.py.
  • Run the prediction script on one of the downloaded flower images by running this command:
python predict.py data/flowers-102/jpg/image_00001.jpg

You should see a list of 5 flower names with confidence percentages printed in the terminal, followed by a matplotlib window showing the image and bar chart.

Not seeing predictions?

If you see a FileNotFoundError for best_model.pth, make sure you ran train.py in the previous step and that the file exists in your project folder.

If the image path is not found, confirm the dataset was downloaded in Step 2. The images live in data/flowers-102/jpg/.

If the matplotlib window does not appear, the chart is still saved as prediction_result.png in your project folder. Open it in your file browser to verify.

help me debug my predict.py script

✔️ Awesome, I've got everything!

Great. Double check you have saved predict.py and confirmed predictions print in the terminal.

ⓧ I'd like to double check the full code

Here is the complete predict.py file:

import sys

import matplotlib.pyplot as plt
import torch
import torch.nn as nn
from torchvision import transforms
from torchvision.models import resnet18
from PIL import Image


FLOWER_NAMES = [
    "pink primrose", "hard-leaved pocket orchid", "canterbury bells",
    "sweet pea", "english marigold", "tiger lily", "moon orchid",
    "bird of paradise", "monkshood", "globe thistle", "snapdragon",
    "colt's foot", "king protea", "spear thistle", "yellow iris",
    "globe-flower", "purple coneflower", "peruvian lily", "balloon flower",
    "giant white arum lily", "fire lily", "pincushion flower",
    "fritillary", "red ginger", "grape hyacinth", "corn poppy",
    "prince of wales feathers", "stemless gentian", "artichoke",
    "sweet william", "carnation", "garden phlox", "love in the mist",
    "mexican aster", "alpine sea holly", "ruby-lipped cattleya",
    "cape flower", "great masterwort", "siam tulip", "lenten rose",
    "barbeton daisy", "daffodil", "sword lily", "poinsettia",
    "bolero deep blue", "wallflower", "marigold", "buttercup",
    "oxeye daisy", "common dandelion", "petunia", "wild pansy",
    "primula", "sunflower", "pelargonium", "bishop of llandaff",
    "gaura", "geranium", "orange dahlia", "pink-yellow dahlia",
    "cautleya spicata", "japanese anemone", "black-eyed susan",
    "silverbush", "californian poppy", "osteospermum", "spring crocus",
    "bearded iris", "windflower", "tree poppy", "gazania",
    "azalea", "water lily", "rose", "thorn apple", "morning glory",
    "passion flower", "lotus", "toad lily", "anthurium",
    "frangipani", "clematis", "hibiscus", "columbine", "desert-rose",
    "tree mallow", "magnolia", "cyclamen", "watercress",
    "canna lily", "hippeastrum", "bee balm", "ball moss",
    "foxglove", "bougainvillea", "camellia", "mallow",
    "mexican petunia", "bromelia", "blanket flower",
    "trumpet creeper", "blackberry lily",
]


def load_model(model_path, num_classes=102):
    model = resnet18(weights=None)
    num_ftrs = model.fc.in_features
    model.fc = nn.Linear(num_ftrs, num_classes)
    model.load_state_dict(torch.load(model_path, map_location="cpu"))
    model.eval()
    return model


def predict_image(model, image_path):
    transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225]
        ),
    ])

    image = Image.open(image_path).convert("RGB")
    input_tensor = transform(image).unsqueeze(0)

    with torch.no_grad():
        outputs = model(input_tensor)
        probabilities = torch.softmax(outputs, dim=1)
        top5_prob, top5_idx = torch.topk(probabilities, 5)

    results = []
    for i in range(5):
        idx = top5_idx[0][i].item()
        prob = top5_prob[0][i].item()
        name = FLOWER_NAMES[idx] if idx < len(FLOWER_NAMES) else f"Class {idx}"
        results.append((name, prob))

    return image, results


def display_prediction(image, results, image_path):
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))

    ax1.imshow(image)
    ax1.set_title(f"Input: {image_path}")
    ax1.axis("off")

    names = [r[0] for r in results]
    probs = [r[1] * 100 for r in results]

    ax2.barh(range(5), probs, color="steelblue")
    ax2.set_yticks(range(5))
    ax2.set_yticklabels(names)
    ax2.set_xlabel("Confidence (%)")
    ax2.set_title("Top 5 Predictions")
    ax2.set_xlim(0, 100)

    plt.tight_layout()
    plt.savefig("prediction_result.png", dpi=100, bbox_inches="tight")
    plt.show()
    print("Prediction saved to prediction_result.png")


def main():
    if len(sys.argv) < 2:
        print("Usage: python predict.py <image_path>")
        print("Example: python predict.py data/flowers-102/jpg/image_00001.jpg")
        sys.exit(1)

    image_path = sys.argv[1]
    model = load_model("best_model.pth")
    image, results = predict_image(model, image_path)

    print("\nPredictions:")
    for name, prob in results:
        print(f"  {name}: {prob*100:.1f}%")

    display_prediction(image, results, image_path)


if __name__ == "__main__":
    main()

Secret mission

Gradual Layer Unfreezing

Your frozen-backbone model learned what it could from the new classifier head alone. What if you could unlock deeper layers of ResNet18 and let them adapt to flower-specific features too? Implement a two-phase training strategy used in production fine-tuning workflows to push your accuracy even higher.

Clean Up Your Resources

Clean Up Your Resources

This project runs entirely locally, so there are no ongoing costs or cloud resources to worry about. Decide whether to keep your project files, or delete them entirely.

Resources you used:

  • Project folder on your Desktop (contains train.py, predict.py, best_model.pth, prediction_result.png).
  • Downloaded Flowers102 dataset inside the data/ folder (approximately 350 MB).
  • Python virtual environment with PyTorch, torchvision, scipy, and matplotlib installed.

Keep everything running

No action needed. Choose this if you want to keep experimenting with your flower classifier, try different images, or tweak your training parameters.

  • Your project files, trained model, and downloaded dataset remain on your Desktop.
  • The virtual environment stays ready for future use.
  • No processes are running in the background, so nothing consumes memory or CPU while idle.

Pause - I'll come back to this later

There are no running processes to stop. Your project is already paused by default since nothing runs in the background.

  • The data/ folder takes up the most disk space (approximately 350 MB). If you need to free space now but want to keep your code, delete just the data/ folder. The dataset will re-download automatically next time you run train.py.
  • Your trained model (best_model.pth) is only about 45 MB and is worth keeping so you can run predictions without retraining.

Delete - I don't want to use this again

Remove all project files and the virtual environment to start fresh.

  • If your virtual environment is active, deactivate it first by running:
deactivate
  • Delete the entire project folder from your Desktop. On macOS or Linux, run:
cd ~/Desktop
rm -rf flower-classifier
  • On Windows, run:
cd ~\Desktop
rmdir /s /q flower-classifier

This removes your Python scripts, the virtual environment, the downloaded dataset, and the trained model in one step.

Nice Work!

Nice Work!

Nice work! You just built a flower image classifier powered by transfer learning, trained entirely on your CPU.

You've learned how to:

  • Build a complete PyTorch training pipeline from scratch, including data loading, image transforms, training and validation loops, and model checkpointing.
  • Apply transfer learning with a pretrained ResNet18 model by freezing the backbone and training only the classifier head on 102 flower categories.
  • Run inference on new images with top-5 confidence scores and visualize predictions using matplotlib.
  • Secret Mission: Implemented gradual layer unfreezing with differential learning rates for improved accuracy, a technique used in production fine-tuning workflows.

Ready to quiz yourself?