用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%
由此可见，卷积神经网络的分类效果远好于全连接网络。