pytorch基础: 分类模型的训练框架

• 以多层感知机MLP模型为例

import torch
from torch import nn, optim
from torch.nn import init
import torchvision
from torchvision import transforms
import numpy as np
import matplotlib.pyplot as plt

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 输入输出
num_inputs = 28 * 28
num_outputs = 10

# 超参数
batch_size = 256
num_workers = 0  # Linux支持多进程读取数据
num_epochs = 10
lr = 0.5


class MLP(nn.Module):
    def __init__(self, num_inputs, num_outputs):
        super(MLP, self).__init__()
        self.mlp = nn.Sequential(
            nn.Linear(num_inputs, 256),  # 第1隐含层 256
            nn.ReLU(),
            nn.Linear(256, 100),  # 第2隐含层 100
            nn.ReLU(),
            nn.Linear(100, num_outputs)  # 输出层
        )

    def forward(self, x):
        x = x.view(x.shape[0], -1)
        x = self.mlp(x)
        return x


def evaluate_accuracy(data_iter, net):
    acc_sum, n = 0.0, 0
    for X, y in data_iter:
        X = X.to(device)
        y = y.to(device)
        acc_sum += (net(X).argmax(dim=1) == y).float().sum().item()
        n += y.shape[0]
    return acc_sum / n


def main():
    # 下载数据集
    root = "./Datasets/FashionMNIST"
    mnist_train = torchvision.datasets.FashionMNIST(root=root, train=True, download=True,
                                                    transform=transforms.ToTensor())
    mnist_test = torchvision.datasets.FashionMNIST(root=root, train=False, download=True,
                                                   transform=transforms.ToTensor())
    # 构建小批量数据生成器
    train_iter = torch.utils.data.DataLoader(mnist_train, batch_size=batch_size, shuffle=True, num_workers=num_workers)
    test_iter = torch.utils.data.DataLoader(mnist_test, batch_size=batch_size, shuffle=False, num_workers=num_workers)

    # 构建模型
    net = MLP(num_inputs, num_outputs)
    net = net.to(device)
    print("training on ", device)
    for params in net.parameters():  # 模型参数初始化
        init.normal_(params, mean=0, std=0.01)
    loss = nn.CrossEntropyLoss()  # 损失函数
    optimizer = optim.SGD(net.parameters(), lr=lr)
    # optimizer = optim.Adam(net.parameters(), lr=lr)

    # 训练模型
    loss_list, train_acc_list, test_acc_list = [], [], []
    for epoch in range(num_epochs):
        train_l_sum, train_acc_sum, n = 0.0, 0.0, 0
        for X, y in train_iter:
            X = X.to(device)
            y = y.to(device)
            y_hat = net(X)
            l = loss(y_hat, y).sum()
            # 梯度清零
            optimizer.zero_grad()
            # 更新梯度
            l.backward()
            optimizer.step()
            # 计算指标
            train_l_sum += l.item()
            train_acc_sum += (y_hat.argmax(dim=1) == y).sum().item()
            n += y.shape[0]
        # 周期评估模型
        test_acc = evaluate_accuracy(test_iter, net)
        loss_list.append(train_l_sum / n)
        train_acc_list.append(train_acc_sum / n)
        test_acc_list.append(test_acc)
        print('epoch %d, loss %.4f, train acc %.3f, test acc %.3f'
              % (epoch + 1, train_l_sum / n, train_acc_sum / n, test_acc))

    # 可视化
    idx = range(1, num_epochs + 1)
    plt.figure(figsize=(12, 4))
    plt.subplot(121)
    plt.plot(idx, loss_list, 'o-')
    plt.grid(), plt.xlabel("epoch"), plt.ylabel("loss"), plt.xticks(range(min(idx), max(idx) + 1, 1))
    plt.subplot(122)
    plt.plot(idx, train_acc_list, 'ro-', label="train accuracy")
    plt.plot(idx, test_acc_list, 'bo-', label="test accuracy")
    plt.xlabel("epoch"), plt.ylabel("accuracy"), plt.legend(loc="best")
    plt.xticks(range(min(idx), max(idx) + 1, 1))
    # plt.yticks(np.arange(0., 1.1, 0.1))
    plt.grid(), plt.ylim([0, 1.1])
    plt.show()


if __name__ == '__main__':
    main()

• 训练结果