【深度学习】VGG16模型训练（CIFAR-10数据集）

VGG16 是一种经典的卷积神经网络（CNN）架构，由牛津大学的 Visual Geometry Group（VGG）提出，特别是在 2014 年的 ImageNet 挑战中表现出色。VGG16 具有简洁而有效的结构，广泛应用于图像分类任务。CIFAR-10（Canadian Institute For Advanced Research 10）是一个常用的小型图像分类数据集，用于训练和测试机器学习模型。它包含了 10 个不同类别的图像，每个类别有 6,000 张图像。

本文使用CIFAR-10数据集，训练VGG16模型，学习实践神经网络模型训练的过程和原理。

VGG16_CIFAR-10 训练测试代码

代码实现了一个基于 VGG16 模型的图像分类任务，使用 CIFAR-10 数据集进行训练和测试

1. 导入必要的库

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

• torch 是 PyTorch 库，提供了各种用于构建和训练神经网络的功能。
• torchvision 是一个常用的图像处理库，提供了预训练模型和数据集等。
• models：从 torchvision 中导入模型，VGG16 就是从这里加载的。
• datasets：用于加载常见的数据集（如 CIFAR-10、ImageNet 等）。
• transforms：用于定义图像预处理操作。
• DataLoader：帮助将数据集分批次加载，用于训练和评估模型。
• time：用于计算每个 epoch 的时间。

2. 设置线程数

torch.set_num_threads(4)

• 设置每个进程使用的线程数为 4，这有助于在多核 CPU 上提高并行性能，尤其在数据加载和训练时。

3. 数据预处理

transform = transforms.Compose([
    transforms.Resize(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

• transforms.Resize(224)：将图像调整为 224x224 的大小，适应 VGG16 模型的输入要求。
• transforms.ToTensor()：将图像转换为 Tensor 格式，并将像素值从 [0, 255] 范围转换到 [0.0, 1.0]。
• transforms.Normalize(mean, std)：对图像进行标准化，将每个通道（RGB）的均值和标准差标准化。这里使用的是在 ImageNet 数据集上计算的均值和标准差，常用于大多数预训练模型（如 VGG16）。

4. 加载 CIFAR-10 数据集

train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
test_dataset = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)

train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True, num_workers=0)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False, num_workers=0)

• datasets.CIFAR10：加载 CIFAR-10 数据集，包含 10 类图像，每类 6,000 张图像。train=True 表示加载为训练集，train=False 表示加载为测试集。
• transform=transform：应用前面定义的图像预处理。
• DataLoader：将数据集分批加载，batch_size=64 表示每次加载 64 张图像。shuffle=True 表示训练集每个 epoch 重新打乱数据，num_workers=0 表示数据加载使用主线程（你可以根据机器的 CPU 核数调整这个值以提高性能）。

5. 加载 VGG16 模型

model = models.vgg16(weights=None)
model.classifier[6] = nn.Linear(model.classifier[6].in_features, 10)

• models.vgg16(weights=None)：加载没有预训练权重的 VGG16 模型。weights=None 表示不加载 ImageNet 上的预训练权重。
• model.classifier[6]：VGG16 的最后一层是一个全连接层，model.classifier[6] 就是最后一层，默认输出 1000 类（ImageNet 上的类别数）。这里将其修改为输出 10 类，以适应 CIFAR-10 数据集。

6. 设置设备（GPU 或 CPU）

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

• 通过 torch.cuda.is_available() 检查是否有可用的 GPU，如果有则使用 GPU，否则使用 CPU。
• model.to(device)：将模型移到指定设备（GPU 或 CPU）。

7. 设置损失函数和优化器

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

• criterion = nn.CrossEntropyLoss()：使用交叉熵损失函数，这通常用于多分类任务。
• optimizer = optim.Adam(model.parameters(), lr=0.001)：使用 Adam 优化器，学习率为 0.001。Adam 是一种自适应学习率优化算法，通常能较好地处理大规模数据集。

8. 训练函数

def train(model, train_loader, criterion, optimizer, num_epochs=10):
    model.train()
    epoch_times = []
    
    for epoch in range(num_epochs):
        start_time = time.time()
        
        running_loss = 0.0
        correct = 0
        total = 0
        
        for i, (inputs, labels) in enumerate(train_loader):
            inputs, labels = inputs.to(device), labels.to(device)
            
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            
            running_loss += loss.item()
            _, predicted = outputs.max(1)
            total += labels.size(0)
            correct += predicted.eq(labels).sum().item()
            
            # 每个 step 打印一次
            if (i + 1) % 10 == 0:
                step_loss = running_loss / (i + 1)
                step_acc = 100. * correct / total
                print(f'Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{len(train_loader)}], Loss: {step_loss:.4f}, Accuracy: {step_acc:.2f}%')

        epoch_loss = running_loss / len(train_loader)
        epoch_acc = 100. * correct / total
        
        epoch_end_time = time.time()
        epoch_duration = epoch_end_time - start_time
        epoch_times.append(epoch_duration)
        
        print("*" * 50)
        print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {epoch_loss:.4f}, Accuracy: {epoch_acc:.2f}%, Time: {epoch_duration:.2f}s')
        print("*" * 50)
        
    avg_epoch_time = sum(epoch_times) / num_epochs
    print(f'\nAverage Epoch Time: {avg_epoch_time:.2f}s')

• model.train()：将模型设置为训练模式，启用 Dropout 和 BatchNorm 等训练时特有的行为。 丢弃(Dropout)是指在训练的过程中，隐藏层的神经元会有一定数量的被丢弃掉，可以防止网络的过度拟合；BatchNorm 是对每一层的输入进行标准化，使其在训练过程中保持均值为 0，方差为 1，从而避免不同层的输入数据分布发生变化，加速收敛。
• optimizer.zero_grad()：清空之前的梯度，以防止它们在计算时累加。在 PyTorch 等深度学习框架中，梯度是累加的，在每次进行反向传播时，新的梯度会加到之前计算出的梯度上，所以需要手动清除。
• loss.backward()：反向传播，计算梯度。反向传播（backpropagation）是计算梯度的过程，在每次计算损失（loss）并执行反向传播时，PyTorch 会自动计算出每个参数的梯度，并将它们累积到 .grad 属性中。
• optimizer.step()：更新模型的参数。
• 计算并打印每个 batch 的损失和准确率，输出每个 epoch 的平均损失、准确率和所用时间。

9. 测试函数

def test(model, test_loader):
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for inputs, labels in test_loader:
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            _, predicted = outputs.max(1)
            total += labels.size(0)
            correct += predicted.eq(labels).sum().item()
    
    accuracy = 100. * correct / total
    print(f'Test Accuracy: {accuracy:.2f}%')

• model.eval()：将模型设置为评估模式，禁用 Dropout 和 BatchNorm 等训练时特有的行为。model.train() 会激活训练模式，启用 Dropout 层（即随机丢弃一部分神经元）以及使用当前 mini-batch 的均值和方差来标准化数据的 BatchNorm 层。（实测这两个模式的训练时长差不多）
• with torch.no_grad()：禁用梯度计算，节省内存和计算资源。
• 计算并输出测试集的准确率。

10. 训练和测试模型

train(model, train_loader, criterion, optimizer, num_epochs=10)
test(model, test_loader)

• train()：训练模型 10 个 epoch。
• test()：在训练完成后，对模型进行测试并输出准确率。

11. 保存模型

# 保存模型权重
torch.save(model.state_dict(), 'vgg16_model.pth')

# 保存整个模型（包括结构和权重）
torch.save(model, 'vgg16_full_model.pth')

完整代码

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

# 设置每个进程使用的线程数
torch.set_num_threads(4)

# 数据预处理
transform = transforms.Compose([
    transforms.Resize(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

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

train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True, num_workers=0)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False, num_workers=0)

# 直接加载没有预训练权重的 VGG16 模型
model = models.vgg16(weights=None)

# 最后一层全连接层的输出类别数为 10（CIFAR-10 有 10 类）
model.classifier[6] = nn.Linear(model.classifier[6].in_features, 10)

# 使用 GPU 如果可用
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# 设置损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 训练函数
def train(model, train_loader, criterion, optimizer, num_epochs=10):
    model.train()
    epoch_times = []
    
    for epoch in range(num_epochs):
        start_time = time.time()
        
        running_loss = 0.0
        correct = 0
        total = 0
        
        for i, (inputs, labels) in enumerate(train_loader):
            inputs, labels = inputs.to(device), labels.to(device)
            
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            
            running_loss += loss.item()
            _, predicted = outputs.max(1)
            total += labels.size(0)
            correct += predicted.eq(labels).sum().item()
            
            # 每个 step 打印一次
            if (i + 1) % 10 == 0:  # 每 10 个 batch 打印一次
                step_loss = running_loss / (i + 1)
                step_acc = 100. * correct / total
                print(f'Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{len(train_loader)}], Loss: {step_loss:.4f}, Accuracy: {step_acc:.2f}%')

        epoch_loss = running_loss / len(train_loader)
        epoch_acc = 100. * correct / total
        
        epoch_end_time = time.time()
        epoch_duration = epoch_end_time - start_time
        epoch_times.append(epoch_duration)
        
        print("*" * 50)
        print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {epoch_loss:.4f}, Accuracy: {epoch_acc:.2f}%, Time: {epoch_duration:.2f}s')
        print("*" * 50)
        
    avg_epoch_time = sum(epoch_times) / num_epochs
    print(f'\nAverage Epoch Time: {avg_epoch_time:.2f}s')

# 测试函数
def test(model, test_loader):
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for inputs, labels in test_loader:
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            _, predicted = outputs.max(1)
            total += labels.size(0)
            correct += predicted.eq(labels).sum().item()
    
    accuracy = 100. * correct / total
    print(f'Test Accuracy: {accuracy:.2f}%')

# 训练和测试模型
train(model, train_loader, criterion, optimizer, num_epochs=10)
test(model, test_loader)

# 代码用于学习模型训练过程与原理，故不保存模型

运行训练脚本

在 PyCharm 中选择之前创建的 Anaconda 环境：
打开 PyCharm ，创建一个新的项目。到 Python 解释器设置，在先前配置的解释器中中，选择conda环境，选择envs目录下之前配置好的pytorch环境。

在pycharm中直接运行模型训练即可

每训练完一个minibatch，会输出损失值和正确率，训练完一个epoch会输出训练时间

训练过程中，GPU内存的占用即VGG模型的大小，这里是10.8GB