模型训练与迁移学习

from torch.amp import autocast, GradScaler

scaler = GradScaler("cuda")
for x, y in loader:
    optimizer.zero_grad(set_to_none=True)
    with autocast("cuda", dtype=torch.float16):
        loss = criterion(model(x), y)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

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

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

train_tf = transforms.Compose([
    transforms.RandomResizedCrop(224, scale=(0.7, 1.0)),
    transforms.RandomHorizontalFlip(),
    transforms.ColorJitter(0.2, 0.2, 0.2),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])
val_tf = 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_ds = datasets.ImageFolder("data/train", train_tf)
val_ds = datasets.ImageFolder("data/val", val_tf)
train_loader = DataLoader(train_ds, batch_size=32, shuffle=True, num_workers=4, pin_memory=True)
val_loader = DataLoader(val_ds, batch_size=64, num_workers=4)

model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)
model.fc = nn.Linear(model.fc.in_features, len(train_ds.classes))
model = model.to(device)

criterion = nn.CrossEntropyLoss(label_smoothing=0.1)
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.05)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=20)

for epoch in range(20):
    model.train()
    for x, y in train_loader:
        x, y = x.to(device), y.to(device)
        optimizer.zero_grad(set_to_none=True)
        loss = criterion(model(x), y)
        loss.backward()
        optimizer.step()
    scheduler.step()

    model.eval()
    correct = total = 0
    with torch.inference_mode():
        for x, y in val_loader:
            x, y = x.to(device), y.to(device)
            pred = model(x).argmax(dim=1)
            correct += int((pred == y).sum())
            total += y.numel()
    print(f"epoch {epoch}: val_acc={correct / total:.4f}")

import torch
from torchvision import models

model = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V2)

# 冻结骨干的前三个 stage，只训 stage4 与分类头
for name, param in model.named_parameters():
    if name.startswith(("conv1", "bn1", "layer1", "layer2")):
        param.requires_grad = False

head_params = list(model.fc.parameters())
backbone_params = [p for n, p in model.named_parameters()
                   if p.requires_grad and not n.startswith("fc")]

optimizer = torch.optim.SGD([
    {"params": backbone_params, "lr": 1e-4},   # 深层用小学习率
    {"params": head_params, "lr": 1e-3},       # 分类头用大学习率
], momentum=0.9, weight_decay=1e-4)

scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=30)
print("trainable tensors:", sum(1 for p in model.parameters() if p.requires_grad))

import random
import numpy as np
import torch

def set_seed(seed=42):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

set_seed(42)
print("seed set")

import torch

def lr_finder(model, loader, criterion, optimizer, device, start=1e-6, end=1.0, steps=100):
    lrs, losses = [], []
    mult = (end / start) ** (1 / steps)
    lr = start
    for i, (x, y) in enumerate(loader):
        if i >= steps:
            break
        for g in optimizer.param_groups:
            g["lr"] = lr
        x, y = x.to(device), y.to(device)
        optimizer.zero_grad(set_to_none=True)
        loss = criterion(model(x), y)
        loss.backward()
        optimizer.step()
        lrs.append(lr)
        losses.append(float(loss))
        lr *= mult
    # 取损失下降最快处的学习率作为上限参考
    return lrs, losses

# lrs, losses = lr_finder(model, train_loader, criterion, optimizer, device)
print("lr finder ready")

import torch
from torchvision import models

model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1).eval()
dummy = torch.randn(1, 3, 224, 224)
with torch.inference_mode():
    torch.onnx.export(
        model, dummy, "resnet18.onnx",
        input_names=["input"], output_names=["logits"],
        dynamic_axes={"input": {0: "batch"}, "logits": {0: "batch"}},
        opset_version=12,
    )
print("exported")