用Pytorch定义并训练一个简单的卷积神经网络

用Pytorch定义并训练一个简单的全连接网络完整展示了使用PyTorch定义模型，载入数据集，训练模型并评估模型的全流程，本文将介绍用Pytorch定义并训练一个简单的卷积神经网络。

首先，请学习卷积神经网络的基础知识和基础组件

其次，本文针对MNIST数据集定义的卷积神经网络如下：

# Optional：Define CNN
class CNN(nn.Module):
    def __init__(self, in_channels=1, num_classes=10):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels=in_channels, out_channels=8, kernel_size=3, stride=1, padding=1) # same convolution
        self.pool = nn.MaxPool2d(kernel_size=(2,2), stride=(2,2)) # Downsampling by 2
        self.conv2 = nn.Conv2d(in_channels=8, out_channels=16, kernel_size=3, stride=1, padding=1)
        self.fc1 = nn.Linear(16*7*7, num_classes)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.reshape(x.shape[0], -1)
        x = self.fc1(x)
        return x

• 卷积层用于提取空间特征，其参数：
  • in_channels指输入Tensor有几个通道
  • out_channels在有些地方也叫滤波器数量，例如：TensorFlow的Conv2D的参数名就叫filters。每个滤波器都在输入Tensor上提取特征
  • kernel_size：定义滤波器Kernel大小
  • stride(步幅)定义滤波器每次移动的像素个数
  • padding(填充)定义在边界周围用零填充
  • kernel_size, stride和padding让我们可以控制输出体积的空间大小，由下面的公式决定：

卷积尺寸输出公式

所以，代码中kernel_size = 3, stride=1, padding=1现实了输出尺寸与输入尺寸同样大小(same convolution)

• 池化层(Pooling): 在Conv 层之间周期性地插入一个 Pooling 层用于逐步减小特征空间的大小，以减少网络中的参数量和计算量，从而也可以控制过拟合。池化层kernel_size=(2,2), stride=(2,2)实现将长宽缩小一半，如下图所示：

  池化层

• 全连接层(Fully connected)：用于分类(classifier)。

训练过程的代码，基本可以完整复用全连接网络的，完整代码如下所示：

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.utils.data import DataLoader
import torchvision.datasets as datasets
import torchvision.transforms as transforms

# Step1: Define Fully connected network
class NN(nn.Module):
    def __init__(self, num_features, num_classes=10):
        super().__init__()
        self.fc1 =nn.Linear(num_features, 50)
        self.fc2 = nn.Linear(50, num_classes)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# Optional：Define CNN
class CNN(nn.Module):
    def __init__(self, in_channels=1, num_classes=10):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels=in_channels, out_channels=8, kernel_size=3, stride=1, padding=1) # same convolution
        self.pool = nn.MaxPool2d(kernel_size=(2,2), stride=(2,2)) # Downsampling by 2
        self.conv2 = nn.Conv2d(in_channels=8, out_channels=16, kernel_size=3, stride=1, padding=1)
        self.fc1 = nn.Linear(16*7*7, num_classes)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.reshape(x.shape[0], -1)
        x = self.fc1(x)
        return x

# Set device & Hyperparameters
device = "cuda" if torch.cuda.is_available() else "cpu"

num_features = 784
num_classes = 10
learning_rate = 1e-3
batch_size = 64
num_epochs = 3

# Step2: Load data
train_dataset = datasets.MNIST(root="dataset/", train=True, transform=transforms.ToTensor(), download=True)
train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

test_dataset = datasets.MNIST(root="dataset/", train=False, transform=transforms.ToTensor(), download=True)
test_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=False)

# Step3: Initialize network
# model = NN(num_features, num_classes).to(device)
model = CNN().to(device)

# Step4: define Loss and optimizer
loss_fn = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

# Step5: Train Network
for epoch in range(num_epochs):
    losses=[]
    for batch_idx, (data, targets) in enumerate(train_dataloader):
        data = data.to(device=device)
        targets = targets.to(device=device)
        # data = data.reshape(data.shape[0], -1)

        # forward
        preds = model(data)
        loss = loss_fn(preds, targets)
        losses.append(loss)
        
        # backward
        optimizer.zero_grad()
        loss.backward()

        # GSD
        optimizer.step()
    print(f"Epoch:{epoch}, loss is {sum(losses)/len(losses)}.")

# Step6: Chekc accuracy on test dataset
num_correct = 0
num_samples = 0
model.eval()

with torch.no_grad():
    for data, targets in test_dataloader:
        data = data.to(device)
        targets = targets.to(device)
        # data = data.reshape(data.shape[0], -1)

        preds = model(data)
        _, results = preds.max(1)
        # print(preds.shape, results.shape, targets.shape)
        num_correct += (results == targets).sum()
        num_samples += results.size(0)
    print(f"The accuracy on test dataset is : {float(num_correct)/float(num_samples)*100:.2f}%")

运行结果如下：

Epoch:0, loss is 0.36820390820503235.
Epoch:1, loss is 0.11146386712789536.
Epoch:2, loss is 0.07877714186906815.
The accuracy on test dataset is : 98.09%

由此可见，卷积神经网络的分类效果远好于全连接网络。