当前位置：首页 > news >正文

介绍如何开发一个小样本增量学习程序

news 来源：原创 2025/5/30 3:58:04

下面将为你详细介绍如何开发一个小样本增量学习程序。我们将使用 PyTorch 框架，并以一个简单的图像分类任务为例进行说明。

整体思路

数据集准备：将数据集划分为旧类和新类，模拟增量学习的场景。
模型定义：定义一个简单的卷积神经网络作为分类器。
旧类训练：在旧类数据集上训练模型。
增量学习：使用小样本新类数据对模型进行增量训练，采用一些技术（如知识蒸馏）来缓解新类损失度高的问题。
测试：在新类和旧类数据集上测试模型的性能。

代码实现

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

# 定义简单的卷积神经网络
class SimpleCNN(nn.Module):
    def __init__(self, num_classes):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, num_classes)

    def forward(self, x):
        x = nn.functional.relu(nn.functional.max_pool2d(self.conv1(x), 2))
        x = nn.functional.relu(nn.functional.max_pool2d(self.conv2(x), 2))
        x = x.view(-1, 320)
        x = nn.functional.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 数据预处理
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

# 加载 MNIST 数据集
train_dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=True)
test_dataset = datasets.MNIST(root='./data', train=False, transform=transform)

# 划分旧类和新类
old_classes = [0, 1, 2, 3, 4]
new_classes = [5, 6, 7, 8, 9]

old_train_indices = [i for i, (_, label) in enumerate(train_dataset) if label in old_classes]
new_train_indices = [i for i, (_, label) in enumerate(train_dataset) if label in new_classes]
old_test_indices = [i for i, (_, label) in enumerate(test_dataset) if label in old_classes]
new_test_indices = [i for i, (_, label) in enumerate(test_dataset) if label in new_classes]

old_train_dataset = Subset(train_dataset, old_train_indices)
new_train_dataset = Subset(train_dataset, new_train_indices)
old_test_dataset = Subset(test_dataset, old_test_indices)
new_test_dataset = Subset(test_dataset, new_test_indices)

# 小样本新类数据
small_sample_size = 100
small_sample_indices = new_train_indices[:small_sample_size]
small_sample_dataset = Subset(train_dataset, small_sample_indices)

# 数据加载器
old_train_loader = DataLoader(old_train_dataset, batch_size=64, shuffle=True)
small_sample_loader = DataLoader(small_sample_dataset, batch_size=64, shuffle=True)
old_test_loader = DataLoader(old_test_dataset, batch_size=64, shuffle=False)
new_test_loader = DataLoader(new_test_dataset, batch_size=64, shuffle=False)

# 初始化模型
model = SimpleCNN(num_classes=10)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 旧类训练
num_epochs = 5
for epoch in range(num_epochs):
    model.train()
    for batch_idx, (data, target) in enumerate(old_train_loader):
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
    print(f'Epoch {epoch+1}/{num_epochs}, Loss: {loss.item()}')

# 增量学习
teacher_model = model.clone()  # 克隆旧模型作为教师模型
num_epochs_incremental = 3
for epoch in range(num_epochs_incremental):
    model.train()
    for batch_idx, (data, target) in enumerate(small_sample_loader):
        optimizer.zero_grad()
        output = model(data)
        teacher_output = teacher_model(data)

        # 知识蒸馏损失
        distillation_loss = nn.KLDivLoss()(nn.functional.log_softmax(output / 2.0, dim=1),
                                           nn.functional.softmax(teacher_output / 2.0, dim=1))
        classification_loss = criterion(output, target)
        loss = classification_loss + distillation_loss

        loss.backward()
        optimizer.step()
    print(f'Incremental Epoch {epoch+1}/{num_epochs_incremental}, Loss: {loss.item()}')

# 测试
def test(model, test_loader):
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for data, target in test_loader:
            output = model(data)
            _, predicted = torch.max(output.data, 1)
            total += target.size(0)
            correct += (predicted == target).sum().item()
    accuracy = 100 * correct / total
    return accuracy

old_accuracy = test(model, old_test_loader)
new_accuracy = test(model, new_test_loader)
print(f'Old class accuracy: {old_accuracy}%')
print(f'New class accuracy: {new_accuracy}%')