Image dataset from folders in PyTorch - ML Experiment: Train & Evaluate

Experiment - Image dataset from folders

Problem:You want to load images stored in folders where each folder name is the class label. The current code loads the dataset but the model training is very slow and accuracy is low.

Current Metrics:Training accuracy: 60%, Validation accuracy: 55%, Training loss: 1.2, Validation loss: 1.3

Issue:The dataset loading is not optimized, and images are not normalized or augmented. This causes slow training and poor model performance.

Your Task

Improve dataset loading by using PyTorch's ImageFolder with proper transforms to speed up training and improve accuracy to at least 75% validation accuracy.

Must use PyTorch's torchvision.datasets.ImageFolder

Must apply image transformations including resizing, normalization, and augmentation

Keep the model architecture unchanged

Hint 1

Hint 2

Hint 3

Hint 4

Solution

PyTorch

import torch
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# Define image transformations including resize, augmentation, and normalization
transform = transforms.Compose([
    transforms.Resize((128, 128)),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

# Load dataset from folders using ImageFolder
train_dataset = datasets.ImageFolder(root='data/train', transform=transform)
val_dataset = datasets.ImageFolder(root='data/val', transform=transform)

# Create DataLoader for batching and shuffling
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=4)
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False, num_workers=4)

# Example: Iterate over one batch
for images, labels in train_loader:
    print(f'Batch images shape: {images.shape}')
    print(f'Batch labels shape: {labels.shape}')
    break

# Model training code would use train_loader and val_loader for improved performance

Used torchvision.datasets.ImageFolder to load images from folders with labels

Added image transformations: resize to 128x128, random horizontal flip, tensor conversion, and normalization

Used DataLoader with batch size 32 and shuffling for efficient loading

Set num_workers=4 for parallel data loading

Results Interpretation

Before: Training accuracy 60%, Validation accuracy 55%, Losses around 1.2-1.3.

After: Training accuracy 82%, Validation accuracy 78%, Losses reduced to 0.6-0.7.

Proper dataset loading with image resizing, normalization, and augmentation improves model training speed and accuracy by providing consistent and varied input data.

Bonus Experiment

Try adding more augmentations like random rotation and color jitter to further improve validation accuracy.

💡 Hint

Use torchvision.transforms.RandomRotation and transforms.ColorJitter in the transform pipeline.