Implementing Image Classification with ResNet-18 on Custom Datasets in PyTorch

Assume your data is structured as folows:

data/
├── dog/
│   ├── dog001.jpg
│   └── dog002.jpg
└── cat/
    ├── cat001.jpg
    └── cat002.jpg

First, import necessary modules and prepare the dataset:

import torch
import torchvision
from torchvision import transforms
import torch.nn as nn
import torch.optim as optim

base_path = './data'

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

full_dataset = torchvision.datasets.ImageFolder(root=base_path, transform=preprocess)

from torch.utils.data import random_split

train_count = int(0.8 * len(full_dataset))
eval_count = len(full_dataset) - train_count
train_set, eval_set = random_split(full_dataset, [train_count, eval_count])

batch = 32
train_loader = torch.utils.data.DataLoader(train_set, batch_size=batch, shuffle=True)
eval_loader = torch.utils.data.DataLoader(eval_set, batch_size=batch, shuffle=False)

Initialize the ResNet-18 model and modify its final layer:

net = torchvision.models.resnet18(pretrained=True)
category_count = len(full_dataset.classes)
net.fc = nn.Linear(net.fc.in_features, category_count)

loss_fn = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

Train the model:

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
net.to(device)

epochs = 10
for epoch_idx in range(epochs):
    net.train()
    cumulative_loss = 0.0
    for inputs, targets in train_loader:
        inputs, targets = inputs.to(device), targets.to(device)
        
        optimizer.zero_grad()
        predictions = net(inputs)
        loss_val = loss_fn(predictions, targets)
        loss_val.backward()
        optimizer.step()
        
        cumulative_loss += loss_val.item()
    
    avg_loss = cumulative_loss / len(train_loader)
    print(f"Epoch {epoch_idx + 1}/{epochs}, Average Loss: {avg_loss:.4f}")

Evaluate the trrained model:

net.eval()
correct_predictions = 0
total_examples = 0

with torch.no_grad():
    for inputs, targets in eval_loader:
        inputs, targets = inputs.to(device), targets.to(device)
        outputs = net(inputs)
        _, predicted_labels = torch.max(outputs, 1)
        total_examples += targets.size(0)
        correct_predictions += (predicted_labels == targets).sum().item()

accuracy = correct_predictions / total_examples
print(f"Evaluation Accuracy: {accuracy * 100:.2f}%")

torch.save(net, "trained_resnet.pth")

For inference on a new image:

from PIL import Image
import torch.nn.functional as F

class_labels = ('cat', 'dog')

model = torch.load('trained_resnet.pth')
model.eval()

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

sample_img = Image.open("./sample_dog.jpg")
sample_tensor = img_transform(sample_img).unsqueeze(0)

with torch.no_grad():
    model_output = model(sample_tensor)
    probabilities = F.softmax(model_output, dim=1)
    _, predicted_idx = torch.max(model_output, 1)
    predicted_class = class_labels[predicted_idx.item()]

print(f"Predicted class: {predicted_class}")
print(f"Class probabilities: {probabilities}")

PyTorch's torcvhision library offers various pre-trained models, such as AlexNet, VGG variants (e.g., VGG11, VGG16 with batch normalization), and ResNet architectures (ResNet-18, ResNet-34, ResNet-50, ResNet-101, ResNet-152). Each can be initialized with pretrained weights from ImageNet by setting pretrained=True.

Tags: pytorch ResNet image classification Custom Dataset Deep Learning

Posted on Tue, 20 Oct 2026 16:53:11 +0000 by salasilm