利用PyTorch实现的深度学习解决MNIST数据集识别代码，并利用GPU训练
深度学习网络一般分为4个部分：
- 数据集的准备和处理
- 定义网络模型
- 定义损失函数和优化器
- 训练和测试
import torch
import torch.nn as nn
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
# 1 data
batch_size = 64 # 批处理的大小
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]) # 先把数据转换为Tensor，再进行归一化
train_dataset = datasets.MNIST(root='../dataset/mnist/', train=True, download=True, transform=transform) # 下载数据集
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True) # 把数据集放到DataLoader中
test_dataset = datasets.MNIST(root='../dataset/mnist/', train=False, download=True, transform=transform)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
# 2 model
# 把残差块定义成一个类，方便代码复用
class ResidualBlock(nn.Module):
def __init__(self, channels):
super(ResidualBlock, self).__init__()
self.channels = channels # 通道数
self.conv1 = nn.Conv2d(channels, channels, kernel_size=3, padding='same') # 3*3的卷积核，padding=same表示卷积后的大小不变
self.conv2 = nn.Conv2d(channels, channels, kernel_size=3, padding='same')
self.relu = nn.ReLU()
def forward(self, x):
y = self.relu(self.conv1(x))
y = self.conv2(y)
return self.relu(x + y)
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(1, 16, kernel_size=5) # 1表示输入的通道数，16表示输出的通道数，kernel_size表示卷积核的大小
self.conv2 = nn.Conv2d(16, 32, kernel_size=5)
self.res1 = ResidualBlock(16) # 16表示输入的通道数，残差网络的输入和输出通道数相同
self.res2 = ResidualBlock(32)
self.relu = nn.ReLU()
self.mp = nn.MaxPool2d(2) # 2表示池化核的大小
self.fc = nn.Linear(512, 10)
def forward(self, x):
in_size = x.size(0)
x = self.mp(self.relu(self.conv1(x))) # 先卷积，再激活，再池化
x = self.res1(x)
x = self.mp(self.relu(self.conv2(x)))
x = self.res2(x)
x = x.view(in_size, -1) # 把x拉成一维的向量，-1表示自适应
x = self.fc(x)
return x
model = Net()
device = torch.device("cuda:0"if torch.cuda.is_available() else "cpu")
model.to(device) # 把模型放到GPU上
# 3 loss and optimizer
loss = nn.CrossEntropyLoss() # 交叉熵损失函数
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.5) # momentum表示动量，作用是加快收敛速度
# 4 training and testing
def train(t_epoch):
running_loss = 0.0
for batch_index, data in enumerate(train_loader, 0): # enumerate函数可以把一个可遍历的数据对象组合成一个索引序列，同时列出数据和数据下标，0表示从0开始计数
inputs, outputs = data
inputs, outputs = inputs.to(device), outputs.to(device) # 把数据放到GPU上
optimizer.zero_grad()
# forward + backward + update
y_hat = model(inputs)
t_loss = loss(y_hat, outputs)
t_loss = t_loss.to(device) # 把损失函数放到GPU上
t_loss.backward()
optimizer.step()
running_loss += t_loss.item()
if batch_index % 300 == 299:
print('[%d, %5d] loss: %.3f' % (t_epoch + 1, batch_index + 1, running_loss / 300))
running_loss = 0.0
def test():
correct = 0
total = 0
with torch.no_grad():
for data in test_loader:
images, labels = data
images, labels = images.to(device), labels.to(device) # 把数据放到GPU上
outputs = model(images)
predicted = torch.argmax(outputs.data, dim=1) # torch.argmax函数可以返回每一行中最大值的索引
total += labels.size(0)
correct += (predicted == labels).sum().item()
print('[%d / %d]' % (correct, total))
print('Accuracy on test set: %d %%' % (100 * correct / total))
if __name__ == '__main__':
for epoch in range(10):
train(epoch)
test()
print('Program finished')
print('Using device: ', device)
print('Using GPU: ', torch.cuda.get_device_name()) # 查看GPU的型号