Pytorch实现简单神经网络的基本模板

神经网络这么火的当下，不学一学会失去很多机会。但网上的资料鱼龙混杂，因此整理自己的学习过程是很有意义的

结构总览

  • 数据预处理
  • 加载数据集
  • 构建模型
  • 训练模型
  • 评估模型

数据预处理

这一步的目的是将数据处理成模型更容易识别的形式，例如裁切、张量化、归一化

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

上述代码中，调用 transforms.Compose，定义了名叫 transform 的转换器
第一行表示将图片裁切为特定大小，便于规范图片格式
第二行将图片张量化，使之成为更易于训练的格式
第三行为归一化：
mean：表示每个通道的均值，是一个长度等于通道数的序列。
std：表示每个通道的标准差，是一个长度等于通道数的序列。
inplace：布尔值，默认是 False。若设为 True，就会原地进行归一化操作。

以上是比较常用的变换操作，当然还有许多其他的变换，如图像翻转和旋转类、图像颜色调整类、图像增强类等，具体根据实际需求选择

加载数据集

这一步的目的是为训练做准备，将数据集导入并合理分配

1
2
3
4
5
trainset = datasets.ImageFolder(root=train_path, transform=transform)
testset = datasets.ImageFolder(root=test_path, transform=transform)

train_loader = DataLoader(trainset, batch_size=64, shuffle=True)
test_loader = DataLoader(testset, batch_size=64, shuffle=True)

这里使用了 datasets.ImageFolder，该类的作用是根据文件夹结构加载图像数据集。例如文件夹A中有B、C两个子文件夹，子文件夹中存储了相应的图片，调用这个类后模型会自动将图片与文件夹名对应，文件夹名即作为图片的标签。

DataLoader 主要功能是将数据集包装成一个可迭代对象，按批次加载数据，便于模型处理。

  • 第一个参数：dataset 加载的数据集对象
  • 第二个参数：batch_size 每个批次加载的数据样本数量 过大会爆显存，过小 则训练慢
  • 第三个参数：shuffle 默认为False.若为True，则打乱数据集的顺序。

还有其他的参数，不过不太常用。

构建模型

核心步骤，用于构建神经网络的具体形式

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, padding=1)
        # 定义第一个卷积层 参数分别是 输入通道数 输出通道数 卷积核大小 填充大小
        self.relu1 = nn.ReLU()
        # 定义第一个 ReLU 激活函数层，用于引入非线性
        self.pool1 = nn.MaxPool2d(2)
        # 定义第一个池化层 类型为最大池化 池化核大小为 2
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
        self.relu2 = nn.ReLU()
        self.pool2 = nn.MaxPool2d(2)
        self.fc1 = nn.Linear(32 * 56 * 56, 128)
        #定义第一个全连接层 参数为 输入特征数 输出特征数
        self.relu3 = nn.ReLU()
        self.fc2 = nn.Linear(128, len(trainset.classes))

    # 前向传播
    def forward(self, x):
        x = self.pool1(self.relu1(self.conv1(x)))
        # 卷积-->激活-->池化
        x = self.pool2(self.relu2(self.conv2(x)))
        x = x.view(-1, 32 * 56 * 56)
        # 将卷积层输出的多维特征图展平为一维向量，以便输入到全连接层。
        x = self.relu3(self.fc1(x))
        # 经过全连接层以及第三个激活函数
        x = self.fc2(x)
        # 经过最后一个全连接层
        return x

上述代码定义了一个简单的卷积神经网络，继承自nn.Module，拥有两个卷积层、两个池化层、三个激活函数层、两个全连接层。
具体的架构需要按照实际需求详细设计

训练模型

这一步实现用加载的数据集对定义好的神经网络类进行训练，一般写在主函数内。
先定义一些基本的类

1
2
3
4
5
6
7
8
model = SimpleCNN()
# 实例化模型类
criterion = nn.CrossEntropyLoss()
# 损失器
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.5)
# 优化器
num_epochs = 10
# 训练轮数

循环迭代训练
总体步骤为 传数据–>算损失–>更新参数

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
for epoch in range(num_epochs):
        model.train()
        # 设置模型为训练模式
        total_loss = 0
        # 总损失

        # 训练模型
        for images, labels in train_loader:
            optimizer.zero_grad()
            # 清空优化器梯度缓存，避免梯度累积从而无法收敛
            outputs = model(images)
            # 将数据传入模型进行训练
            loss = criterion(outputs, labels)
            # 计算损失率
            loss.backward()
            # 反向传播
            optimizer.step()
            # 更新参数

            total_loss += loss.item()
        print(f"Epoch {epoch + 1}/{num_epochs}, Loss: {total_loss:.4f}")

        #----------------------------------------------------
        # 评估模型
        model.eval()
        # 模型设置为评估模式
        correct = 0
        total = 0
        with torch.no_grad():# 避免不必要的梯度计算
            for images, labels in test_loader:
                outputs = model(images)
                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()

        print(f"Test Accuracy: {100 * correct / total:.2f}%")

评估模型

用训练好的模型，在测试集上测试模型性能
可以用下面代码保存训练好的模型

1
torch.save(model, 'result.pth')