Pytorch中的VGG实现修改最后一层FC

下面是PyTorch中修改VGG网络最后一层全连接层的攻略：

步骤一：导入相关库

首先需要导入相关的PyTorch库，主要包括：

• torch：PyTorch的核心库；
• torchvision：PyTorch的图像处理库，提供了很多常用的卷积神经网络的实现，包括VGG等；
• nn：PyTorch中的神经网络模块，用于构建神经网络模型。

步骤二：定义VGG模型

导入VGG网络并定义网络架构，这里以VGG16为例：

import torch
import torchvision.models as models
import torch.nn as nn

vgg16 = models.vgg16(pretrained=True)

此处通过torchvision提供的models模块，调用VGG16的预训练模型。预训练模型参数存储在torch的默认路径中，使用参数pretrained=True即可将模型下载并存储在系统的torch缓存目录中。

步骤三：修改VGG模型中的全连接层

经过上一步的操作，我们就可以得到一个已经训练好的VGG16模型vgg16。但是，VGG网络的最后一层是全连接层，并且输出维度为1000，这并不适合我们训练的任务。因此，我们需要将最后一个全连接层修改为符合任务需求的全连接层。

修改方法如下：

num_features = vgg16.classifier[6].in_features
vgg16.classifier[6] = nn.Linear(num_features, num_classes)

其中，num_features是原模型vgg16中全连接层的输入维度，新建了一个全连接层，将其输入维度设置为num_features，输出维度设置为任务需要的num_classes。这样，我们就完成了VGG16模型中最后一个全连接层的修改。

步骤四：训练修改后的VGG模型

将修改后的VGG16模型作为我们的模型进行训练，具体的训练方法在此不再详述。下面给出两个示例。

示例1：对CIFAR-10数据集进行分类使用

import torch.optim as optim
import torchvision.transforms as transforms
from torchvision.datasets import CIFAR10

# 加载 CIFAR-10 数据集
transform = transforms.Compose([transforms.Resize((224, 224)), transforms.ToTensor()])
trainset = CIFAR10(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=32, shuffle=True, num_workers=2)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(vgg16.parameters(), lr=0.001, momentum=0.9)

# 训练修改后的VGG16模型
for epoch in range(10):
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data
        optimizer.zero_grad()

        outputs = vgg16(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        if i % 2000 == 1999:
            print('[%d, %5d] loss: %.3f' % (epoch + 1, i + 1, running_loss / 2000))
            running_loss = 0.0

print('Finished Training')

示例2：使用自定义数据集训练VGG16模型

import torch.optim as optim
from torch.utils.data import DataLoader
from mydataset import MyDataset

# 加载自定义数据集
trainset = MyDataset(root_dir='./data', transform=transforms.ToTensor())
trainloader = DataLoader(trainset, batch_size=32, shuffle=True, num_workers=2)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(vgg16.parameters(), lr=0.001, momentum=0.9)

# 训练修改后的VGG16模型
for epoch in range(10):
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data
        optimizer.zero_grad()

        outputs = vgg16(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        if i % 2000 == 1999:
            print('[%d, %5d] loss: %.3f' % (epoch + 1, i + 1, running_loss / 2000))
            running_loss = 0.0

print('Finished Training')

在这个示例中，我们自定义了数据集，并使用pytorch中的DataLoader类将自定义的数据集加载到内存中，进行训练。需要注意的是，此处自定义数据集需要自己实现MyDataset类。可以参考如下模板代码：

from torch.utils.data import Dataset

class MyDataset(Dataset):
    def __init__(self, root_dir, transform=None):
        self.root_dir = root_dir
        self.transform = transform
        # 自定义数据集的读取操作

    def __len__(self):
        # 返回数据集的大小
        pass

    def __getitem__(self, idx):
        # 返回一个样本数据
        pass

至此，我们已经完成了修改VGG16网络最后一个全连接层并利用自定义数据集进行训练的整个过程。