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.