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

# 设备与超参数配置
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
batch_size, num_classes, num_epochs = 32, 5, 20
fine_tune_lr = 1e-5  # 微调学习率需偏小，避免破坏预训练特征

# 数据增强（缓解小样本过拟合）与归一化
train_transform = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
    transforms.ToTensor(),
    transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225))
])
val_transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225))
])

# 加载数据集（按类别分文件夹存储）
train_loader = DataLoader(datasets.ImageFolder('./data/animal/train', train_transform),
                          batch_size=batch_size, shuffle=True, num_workers=4)
val_loader = DataLoader(datasets.ImageFolder('./data/animal/val', val_transform),
                        batch_size=batch_size, shuffle=False, num_workers=4)

python
# 加载ResNet50预训练模型，冻结浅层参数
model = models.resnet50(pretrained=True)
for param in model.parameters():
    param.requires_grad = False  # 冻结浅层，保留通用特征

# 修改输出层适配5分类任务，仅训练输出层参数
model.fc = nn.Linear(model.fc.in_features, num_classes)
model = model.to(device)

# 定义损失函数与优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.fc.parameters(), lr=fine_tune_lr)

python
def train_model(model, train_loader, val_loader, criterion, optimizer, num_epochs):
    for epoch in range(num_epochs):
        # 训练阶段
        model.train()
        train_loss = 0.0
        for imgs, lbs in train_loader:
            imgs, lbs = imgs.to(device), lbs.to(device)
            outputs = model(imgs)
            loss = criterion(outputs, lbs)
            
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            train_loss += loss.item() * imgs.size(0)
        
        # 验证阶段
        model.eval()
        val_loss, correct, total = 0.0, 0, 0
        with torch.no_grad():
            for imgs, lbs in val_loader:
                imgs, lbs = imgs.to(device), lbs.to(device)
                outputs = model(imgs)
                val_loss += criterion(outputs, lbs).item() * imgs.size(0)
                _, preds = torch.max(outputs.data, 1)
                total += lbs.size(0)
                correct += (preds == lbs).sum().item()
        
        # 打印训练日志
        train_loss /= len(train_loader.dataset)
        val_loss /= len(val_loader.dataset)
        print(f'Epoch [{epoch+1}/{num_epochs}], Train Loss: {train_loss:.4f}, '
              f'Val Loss: {val_loss:.4f}, Val Acc: {100*correct/total:.2f}%')

# 启动训练
train_model(model, train_loader, val_loader, criterion, optimizer, num_epochs)