简单卷积神经网络Pytorch实现（MNIST数据集分类）

import torch
import torch.utils.data
from torchvision import datasets, transforms
from tqdm import tqdm
# 指定设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print('device:{}'.format(device))
# 定义预处理
transform = transforms.Compose(
    [
        transforms.ToTensor(),
        transforms.Normalize([0.5], [0.5])
    ])
# 导入数据集
dataset_train = datasets.MNIST(root='./data', transform=transform, train=True, download=True)
dataset_test = datasets.MNIST(root='./data', transform=transform, train=False)
# 数据装载
dataloader_train = torch.utils.data.DataLoader(dataset=dataset_train, batch_size=64, shuffle=True)
dataloader_test = torch.utils.data.DataLoader(dataset=dataset_test, batch_size=64, shuffle=False)


# 定义模型
class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = torch.nn.Sequential(
            torch.nn.Conv2d(1, 64, kernel_size=3, stride=1, padding=1),
            torch.nn.ReLU(),
            torch.nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(stride=2, kernel_size=2)
        )

        self.dense = torch.nn.Sequential(
            torch.nn.Linear(14*14*128, 1024),
            torch.nn.ReLU(),
            torch.nn.Dropout(p=0.5),
            torch.nn.Linear(1024, 10)
        )

    def forward(self, x):
        x = self.conv1(x)
        x = x.view(-1, 14*14*128)
        x = self.dense(x)
        return x


model = Net()
model.to(device)
loss_function = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters())

# 训练
for epoch in range(5):
    # 获取该轮损失和准确率
    running_loss = 0.0
    running_accuracy = 0
    # 分批训练
    for data in tqdm(dataloader_train):
        inputs, labels = data
        inputs, labels = inputs.to(device), labels.to(device)
        outputs = model(inputs)
        _, pred = torch.max(outputs, 1)

        optimizer.zero_grad()
        loss = loss_function(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        running_accuracy += (pred == labels).sum().item()
    print('training loss: ', running_loss/len(dataset_train))
    print('training accuracy: ', running_accuracy/len(dataset_train))

# 测试
testing_accuracy = 0
for data in tqdm(dataloader_test):
    inputs, labels = data
    inputs, labels = inputs.to(device), labels.to(device)
    outputs = model(inputs)
    _, pred = torch.max(outputs, 1)

    testing_accuracy += (pred == labels).sum().item()
print('testing accuracy: ', testing_accuracy / len(dataset_test))