CNN好莱坞明星识别

李恒道

数据集https://www.heywhale.com/mw/dataset/63072841b7dcdc54975ad5d0
用VGG-16代替网络

model = models.vgg16(pretrained=True)
for param in model.parameters():
    param.requires_grad = False  # 冻结前面所有层
model.classifier._modules['6'] = nn.Linear(4096, len(classeNames))  # 只训练最后一层
model = model.to(device)

其他什么都不用动，完整代码

import torch
train_acc  = []
import torch
import torch.nn as nn
import torchvision.transforms as transforms
import torchvision
from torchvision import transforms, datasets, models
import torch.nn.functional as F
import os
import PIL
import pathlib
import random

data_dir = './person_photos/'
data_dir = pathlib.Path(data_dir)
total_datadir = './person_photos/'
data_paths = list(data_dir.glob('*'))
classeNames = [str(path).split("\\")[1] for path in data_paths]
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
train_transforms = transforms.Compose([
    transforms.Resize([224, 224]),
    transforms.ToTensor(),
    transforms.Normalize(
        mean=[0.4849, 0.4179, 0.3943],
        std=[0.3205, 0.2929, 0.2896])
])

total_data = datasets.ImageFolder(total_datadir, transform=train_transforms)
train_size = int(0.8 * len(total_data))
test_size = len(total_data) - train_size
# 使用固定随机种子保证每次划分一致
seed = 42
gen = torch.Generator()
gen.manual_seed(seed)
train_dataset, test_dataset = torch.utils.data.random_split(total_data, [train_size, test_size], generator=gen)
print(f"训练集样本数: {len(train_dataset)}")
print(f"测试集样本数: {len(test_dataset)}")

batch_size = 32

def make_dataloaders(train_dataset, test_dataset, batch_size=32, num_workers=0):
    train_dl = torch.utils.data.DataLoader(train_dataset,
                                           batch_size=batch_size,
                                           shuffle=True,
                                           num_workers=num_workers)
    test_dl = torch.utils.data.DataLoader(test_dataset,
                                          batch_size=batch_size,
                                          shuffle=True,
                                          num_workers=num_workers)
    return train_dl, test_dl

# 使用 VGG-16 替代自定义网络，并调整最后的分类层以匹配类别数
num_classes = len(classeNames)
model = models.vgg16(pretrained=True)
for param in model.parameters():
    param.requires_grad = False  # 冻结前面所有层
model.classifier._modules['6'] = nn.Linear(4096, len(classeNames))  # 只训练最后一层
model = model.to(device)

def test(dataloader, model, loss_fn):
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    test_loss, test_acc = 0, 0
    with torch.no_grad():
        for imgs, target in dataloader:
            imgs, target = imgs.to(device), target.to(device)
            target_pred = model(imgs)
            loss = loss_fn(target_pred, target)
            test_loss += loss.item()
            test_acc += (target_pred.argmax(1) == target).type(torch.float).sum().item()

    test_acc /= size
    test_loss /= num_batches
    return test_acc, test_loss

def train(dataloader, model, loss_fn, optimizer):
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    train_loss, train_acc = 0, 0
    for X, y in dataloader:
        X, y = X.to(device), y.to(device)
        pred = model(X)
        loss = loss_fn(pred, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        train_acc += (pred.argmax(1) == y).type(torch.float).sum().item()
        train_loss += loss.item()

    train_acc /= size
    train_loss /= num_batches
    return train_acc, train_loss

train_dl, test_dl = make_dataloaders(train_dataset, test_dataset, batch_size=batch_size, num_workers=0)
batch = next(iter(test_dl))
X, y = batch

print("Shape of X [N, C, H, W]: ", X.shape)
print("Shape of y: ", y.shape, y.dtype)

loss_fn = nn.CrossEntropyLoss()
learn_rate = 1e-3
opt = torch.optim.SGD(model.parameters(), lr=learn_rate, momentum=0.9, weight_decay=1e-4)

# 动态学习率配置：可选择 'cosine', 'step', 'reduce_on_plateau'
scheduler_type = 'cosine'
epochs = 20

# 创建调度器
if scheduler_type == 'cosine':
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=epochs)
elif scheduler_type == 'step':
    scheduler = torch.optim.lr_scheduler.StepLR(opt, step_size=5, gamma=0.5)
elif scheduler_type == 'reduce_on_plateau':
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(opt, mode='min', factor=0.5, patience=3)
else:
    scheduler = None
train_loss = []
train_acc = []
test_loss = []
test_acc = []

for epoch in range(epochs):
    model.train()
    epoch_train_acc, epoch_train_loss = train(train_dl, model, loss_fn, opt)
    model.eval()
    epoch_test_acc, epoch_test_loss = test(test_dl, model, loss_fn)
    train_acc.append(epoch_train_acc)
    train_loss.append(epoch_train_loss)
    test_acc.append(epoch_test_acc)
    test_loss.append(epoch_test_loss)
    # 根据调度器类型更新学习率
    if scheduler is not None:
        if scheduler_type == 'reduce_on_plateau':
            scheduler.step(epoch_test_loss)
        else:
            scheduler.step()

    current_lr = opt.param_groups[0]['lr']
    template = ('Epoch:{:2d}, LR:{:.6f}, Train_acc:{:.1f}%, Train_loss:{:.3f}, Test_acc:{:.1f}%，Test_loss:{:.3f}')
    print(template.format(epoch+1, current_lr, epoch_train_acc*100, epoch_train_loss, epoch_test_acc*100, epoch_test_loss))
torch.save(model.state_dict(), f'best_model_epoch{epoch}.pth')
print('Done')

Epoch:20, LR:0.000000, Train_acc:57.5%, Train_loss:1.328, Test_acc:48.3%，Test_loss:1.613