PyTorch小案例——手写字识别

import torch
import torchvision.datasets as dataset
import torch.utils.data as data_utils
from torchvision import transforms

# 训练集下载
train_data = dataset.MNIST(root='mnist', train=True, transform=transforms.ToTensor(), download=True)
# 测试集下载
test_data = dataset.MNIST(root='mnist', train=False, transform=transforms.ToTensor(), download=False)
# batchsize dataloader可以将数据进行循环读取
train_loader = data_utils.DataLoader(dataset=train_data, batch_size=64, shuffle=True)
test_loader = data_utils.DataLoader(dataset=test_data, batch_size=64, shuffle=True)


class CNN(torch.nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        # 定义卷积层
        self.conv = torch.nn.Sequential(torch.nn.Conv2d(1, 32, kernel_size=5, padding=2), torch.nn.BatchNorm2d(32),
                                        torch.nn.ReLU(),
                                        torch.nn.MaxPool2d(2))
        # 定义网络层
        self.fc = torch.nn.Linear(14 * 14 * 32, 10)

    def forward(self, x):
        out = self.conv(x)
        out = out.view(out.size()[0], -1)
        out = self.fc(out)
        return out


cnn = CNN()
# loss 损失函数，采用交叉熵
loss_func = torch.nn.CrossEntropyLoss()

optimizer = torch.optim.Adam(cnn.parameters(), lr=0.01)

for epoch in range(10):
    for i, (images, labels) in enumerate(train_loader):
        outputs = cnn(images)
        loss = loss_func(outputs, labels)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        print(f'epoch is {epoch + 1}, ite is {len(train_data)}/{64}, loss is {loss.item()}')

# 保存训练好的模型文件
torch.save(cnn, 'mnist.pkl')