Pytorch高级api搭建多层感知机实战

代码

import torch
import torch.nn.functional as F
import torch.optim as optim

from torchvision import datasets, transforms

batch_size = 200
learning_rate = 0.01
epochs = 10

train_loader = torch.utils.data.DataLoader(
    datasets.MNIST('../data',train=True,download=True,
                   transform=transforms.Compose([
                       transforms.ToTensor(),
                       transforms.Normalize((0.1307,),(0.3081,))
                   ]
                   )),
    batch_size=batch_size,shuffle=True,pin_memory=True  # pin_memory=True将数据加载到固定的内存中，以加快数据传输速度
)

test_loader = torch.utils.data.DataLoader(
    datasets.MNIST('../data',train=False,download=True,
                   transform=transforms.Compose([
                       transforms.ToTensor(),
                       transforms.Normalize((0.1307,),(0.3081,))
                   ]
                   )),
    batch_size=batch_size,shuffle=True,pin_memory=True
)


class MLP(torch.nn.Module):
    def __init__(self):
        super(MLP,self).__init__()
        
        # 定义了一个包含多个线性层和LeakyReLU激活函数的前向计算模块
        self.model = torch.nn.Sequential(
            torch.nn.Linear(784,200),
            torch.nn.LeakyReLU(inplace=True),
            torch.nn.Linear(200, 200),
            torch.nn.LeakyReLU(inplace=True),
            torch.nn.Linear(200, 10),
            torch.nn.LeakyReLU(inplace=True),
        )

    def forward(self,x):
        x = self.model(x)

        return x

# 确定使用的设备（GPU或CPU），并将定义的MLP模型移动到该设备上
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
net = MLP().to(device)


optimizer = optim.SGD(net.parameters(),lr=learning_rate)
criteon = torch.nn.CrossEntropyLoss()

for epoch in range(epochs):
    """这是训练过程的主循环。对于每个epoch，循环遍历训练数据加载器中的每个批次。将输入数据和目标标签移动到相同的设备上。通过模型进行前向传播，计算输出logits。然后，使用损失函数计算预测结果和目标标签之间的损失。接下来，执行优化器的相关操作，包括梯度清零、反向传播和参数更新。如果批次索引是100的倍数，则打印训练进度信息。""" 
    for batch_idx,(data,target) in enumerate(train_loader):
        # 将输入数据和目标标签移动到相同的设备上
        data = data.view(-1,28*28).to(device)
        target =target.to(device)

        logits = net(data)
        loss = criteon(logits,target)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if (batch_idx % 100) == 0:
            print('Train Epoch:{} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch,
                batch_idx * len(data),
                len(train_loader.dataset),
                100. * batch_idx / len(train_loader),
                loss.item()
            ))
            
    
    """这段代码用于在训练完成后对模型进行测试。循环遍历测试数据加载器中的每个批次。将输入数据和目标标签移动到设备上，通过模型进行前向传播，计算输出logits。使用损失函数计算预测结果和目标标签之间的损失，并将其累加到test_loss变量中。计算预测的标签，并统计正确的预测数量。最后，计算平均测试损失和准确率，并打印结果。"""
    test_loss = 0
    correct = 0
    for data,target in test_loader:
        data = data.view(-1,28 * 28).to(device)
        target = target.to(device)
        logits = net(data)
        test_loss += criteon(logits,target).item()

        pred = logits.data.max(1)[1]
        correct += pred.eq(target).sum().item()

    test_loss /= len(test_loader.dataset)
    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss,
        correct,
        len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)
    ))