Computer Vision with PyTorch

import torchvision
import torchvision.transforms as T
from torch.utils.data import DataLoader

# Define transforms
train_transform = T.Compose([
    T.RandomResizedCrop(224),
    T.RandomHorizontalFlip(),
    T.ColorJitter(brightness=0.2, contrast=0.2),
    T.ToTensor(),
    T.Normalize(mean=[0.485, 0.456, 0.406],
                std=[0.229, 0.224, 0.225])
])

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

# Load built-in datasets
train_data = torchvision.datasets.CIFAR10(
    root='./data', train=True, download=True, transform=train_transform
)
train_loader = DataLoader(train_data, batch_size=32, shuffle=True, num_workers=4)

# Load custom dataset from folder structure
custom_data = torchvision.datasets.ImageFolder(
    root='data/train', transform=train_transform
)

from torchvision import models

# Load pre-trained models with new weights API
resnet50 = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V2)
efficientnet = models.efficientnet_b0(weights=models.EfficientNet_B0_Weights.DEFAULT)
vit = models.vit_b_16(weights=models.ViT_B_16_Weights.DEFAULT)

# Available models include:
# ResNet (18, 34, 50, 101, 152)
# EfficientNet (B0-B7)
# Vision Transformer (ViT)
# MobileNet V2/V3
# DenseNet, VGG, etc.

import torch
import torch.nn as nn
from torchvision import models

# 1. Load pre-trained model
model = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V2)

# 2. Freeze all layers
for param in model.parameters():
    param.requires_grad = False

# 3. Replace the classification head
num_features = model.fc.in_features
model.fc = nn.Sequential(
    nn.Linear(num_features, 256),
    nn.ReLU(),
    nn.Dropout(0.5),
    nn.Linear(256, 5)  # 5 classes
)

# 4. Train only the new layers
optimizer = torch.optim.Adam(model.fc.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()

# Phase 1: Train the head
model = model.to(device)
for epoch in range(10):
    train_one_epoch(model, train_loader, criterion, optimizer, device)

# Phase 2: Fine-tune last few layers
for param in model.layer4.parameters():
    param.requires_grad = True

optimizer = torch.optim.Adam(model.parameters(), lr=1e-5)  # Lower LR!
for epoch in range(5):
    train_one_epoch(model, train_loader, criterion, optimizer, device)