import torch
from torch import nn
from torchvision import models

# Downloads ImageNet weights on the first run.
model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
model.fc = nn.Linear(model.fc.in_features, 3)
trainable_groups = ["fc"]

for name, parameter in model.named_parameters():
    group = "stem" if name.split(".")[0] in {"conv1", "bn1"} else name.split(".")[0]
    parameter.requires_grad = group in trainable_groups

def set_training_mode():
    # Call at the START of EVERY training epoch.
    model.train()
    # Freeze BN buffers separately from affine weight/bias parameters.
    for module in model.modules():
        if isinstance(module, nn.modules.batchnorm._BatchNorm):
            module.eval()

set_training_mode()
parameters = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.AdamW(parameters, lr=1e-4, weight_decay=1e-2)
print("Total:", sum(p.numel() for p in model.parameters()))
print("Trainable:", sum(p.numel() for p in parameters))
# Use model.eval() and torch.inference_mode() for validation/inference.
# Freezing weights does not remove their forward computation.