Transfer Learning

import torch, torch.nn as nn
from torchvision import models, transforms, datasets
from torch.utils.data import DataLoader

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

# 1. Pre-trained backbone
net = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
for p in net.parameters(): p.requires_grad = False   # freeze
net.fc = nn.Linear(net.fc.in_features, 5)            # 5 class
net = net.to(device)

# 2. Data + transforms
tfm = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize([0.485,0.456,0.406],
                         [0.229,0.224,0.225]),
])
train = datasets.ImageFolder("data/train", transform=tfm)
loader = DataLoader(train, batch_size=32, shuffle=True)

# 3. Loss + optimizer (only fc params)
opt  = torch.optim.Adam(net.fc.parameters(), lr=1e-3)
loss = nn.CrossEntropyLoss()

# 4. Train loop
for epoch in range(5):
    net.train()
    for x, y in loader:
        x, y = x.to(device), y.to(device)
        opt.zero_grad()
        out = net(x); l = loss(out, y); l.backward(); opt.step()
    print(f"Epoch {epoch}: loss {l.item():.3f}")

# Stage 2: শেষ block unfreeze
for p in net.layer4.parameters(): p.requires_grad = True
opt = torch.optim.Adam([
    {"params": net.layer4.parameters(), "lr": 1e-4},
    {"params": net.fc.parameters(),     "lr": 1e-3},
])