Computer Vision

python
Run Inline
import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
# Load CIFAR-10
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])
trainset = torchvision.datasets.CIFAR10(
    root='./data', train=True, download=True, transform=transform
)
trainloader = torch.utils.data.DataLoader(
    trainset, batch_size=32, shuffle=True
)
# Pre-trained ResNet
model = torchvision.models.resnet18(pretrained=True)
# Freeze layers
for param in model.parameters():
    param.requires_grad = False
# Replace classifier
model.fc = nn.Linear(512, 10)