如何增量训练人脸识别模型，无需从头重新训练？

已基于克里斯蒂亚诺·罗纳尔多（Cristiano Ronaldo）和利昂内尔·梅西（Lionel Messi）的图像训练完成一个人脸识别模型，现在希望为模型新增识别人物（如玛丽亚·莎拉波娃（Maria Sharapova）），但不想从头重新训练整个模型，询问是否可以基于新数据集训练并高效合并到现有模型中。

现有训练代码如下：

import torch
import torchvision
from torchvision import datasets, models, transforms
import os

import ssl
ssl._create_default_https_context = ssl._create_unverified_context

data_transforms = {
    'train': transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
    'test': transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
}

data_dir = './new_dataset'
image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x])
                  for x in ['train', 'test']}
dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=4, shuffle=True)
               for x in ['train', 'test']}
class_names = image_datasets['train'].classes

model = models.resnet18(pretrained=True, progress=True)

num_classes = len(class_names)
model.fc = torch.nn.Linear(model.fc.in_features, num_classes)

device = torch.device("cpu")
model = model.to(device)

criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

num_epochs = 10

for epoch in range(num_epochs):
    for inputs, labels in dataloaders['train']:
        inputs = inputs.to(device)
        labels = labels.to(device)

        optimizer.zero_grad()

        outputs = model(inputs)
        loss = criterion(outputs, labels)

        loss.backward()
        optimizer.step()

torch.save(model.state_dict(), 'model.pth')

model.eval()

correct = 0
total = 0

with torch.no_grad():
    for inputs, labels in dataloaders['test']:
        inputs = inputs.to(device)
        labels = labels.to(device)

        outputs = model(inputs)
        _, predicted = torch.max(outputs.data, 1)

        total += labels.size(0)
        correct += (predicted == labels).sum().item()

accuracy = 100 * correct / total
print(f"Accuracy on the test set: {accuracy}%")

原数据集文件夹结构：

new_dataset/
--test/
----cristiano_ronaldo
----lionel_messi
--train/
----cristiano_ronaldo
----lionel_messi

完全可以通过增量训练实现该需求，无需从头训练整个模型。核心思路是保留已训练好的特征提取层权重，仅调整最后一层分类器（或微调少量特征层）适配新类别，具体步骤如下：

1. 调整数据集结构 #

在原数据集的train和test目录下新增玛丽亚·莎拉波娃的图像文件夹：

new_dataset/
--test/
----cristiano_ronaldo
----lionel_messi
----maria_sharapova
--train/
----cristiano_ronaldo
----lionel_messi
----maria_sharapova

2. 加载现有模型并修改分类层 #

加载之前保存的model.pth，将最后一层全连接层（FC）的输出维度从2修改为3（对应3个类别），同时保留原FC层中对应C罗和梅西的权重，仅初始化新增类别的权重：

import torch
import torchvision
from torchvision import datasets, models, transforms
import os

import ssl
ssl._create_default_https_context = ssl._create_unverified_context

# 数据变换保持不变
data_transforms = {
    'train': transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
    'test': transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
}

# 加载更新后的数据集
data_dir = './new_dataset'
image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x])
                  for x in ['train', 'test']}
dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=4, shuffle=True)
               for x in ['train', 'test']}
class_names = image_datasets['train'].classes
num_classes = len(class_names)  # 现在为3

# 加载ResNet18结构，读取已训练好的模型权重
model = models.resnet18(pretrained=False)
original_state_dict = torch.load('model.pth')

# 构建原FC层并加载权重，用于提取原类别参数
original_fc = torch.nn.Linear(model.fc.in_features, 2)
original_fc.load_state_dict({
    'weight': original_state_dict['fc.weight'],
    'bias': original_state_dict['fc.bias']
})

# 创建新FC层，保留原类别权重，初始化新增类别参数
new_fc = torch.nn.Linear(model.fc.in_features, num_classes)
with torch.no_grad():
    new_fc.weight[:2] = original_fc.weight
    new_fc.bias[:2] = original_fc.bias
model.fc = new_fc

# 加载原模型除FC层外的其他参数
del original_state_dict['fc.weight']
del original_state_dict['fc.bias']
model.load_state_dict(original_state_dict, strict=False)

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

3. 冻结特征提取层，仅训练新分类层 #

为了高效训练，冻结ResNet18的所有特征提取层参数，只训练新的FC层：

# 冻结特征提取层参数
for param in model.parameters():
    param.requires_grad = False
# 仅开启FC层参数的梯度更新
for param in model.fc.parameters():
    param.requires_grad = True

# 定义损失函数和优化器（仅优化FC层参数）
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.fc.parameters(), lr=0.001, momentum=0.9)
num_epochs = 10

# 训练过程
for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0
    for inputs, labels in dataloaders['train']:
        inputs = inputs.to(device)
        labels = labels.to(device)

        optimizer.zero_grad()

        outputs = model(inputs)
        loss = criterion(outputs, labels)

        loss.backward()
        optimizer.step()

        running_loss += loss.item() * inputs.size(0)
    
    epoch_loss = running_loss / len(image_datasets['train'])
    print(f'Epoch {epoch+1}/{num_epochs}, Loss: {epoch_loss:.4f}')

# 保存更新后的模型
torch.save(model.state_dict(), 'updated_model.pth')

# 测试模型精度
model.eval()
correct = 0
total = 0

with torch.no_grad():
    for inputs, labels in dataloaders['test']:
        inputs = inputs.to(device)
        labels = labels.to(device)

        outputs = model(inputs)
        _, predicted = torch.max(outputs.data, 1)

        total += labels.size(0)
        correct += (predicted == labels).sum().item()

accuracy = 100 * correct / total
print(f"Accuracy on the test set: {accuracy}%")

4. 可选：微调部分特征层（提升精度） #

如果新数据集图像数量较多，可以解冻ResNet18的最后几个卷积层（如layer4），和FC层一起微调，让模型更好适配新类别的特征：

# 解冻layer4和FC层参数
for param in model.layer4.parameters():
    param.requires_grad = True
for param in model.fc.parameters():
    param.requires_grad = True

# 使用分层学习率微调，避免破坏已训练特征
optimizer = torch.optim.SGD([
    {'params': model.layer4.parameters(), 'lr': 0.0001},
    {'params': model.fc.parameters(), 'lr': 0.001}
], momentum=0.9)

这种方式既保留了原模型对C罗和梅西的识别能力，又高效新增了对莎拉波娃的识别，无需从头训练整个模型。

内容的提问来源于stack exchange，提问作者Sammy