from model import Model import numpy as np import torch from torchvision.datasets import mnist from torch.nn import CrossEntropyLoss from torch.optim import SGD from torch.utils.data import DataLoader from torchvision.transforms import ToTensor if __name__ == '__main__': batch_size = 256 train_dataset = mnist.MNIST(root='./train', train=True, transform=ToTensor()) test_dataset = mnist.MNIST(root='./test', train=False, transform=ToTensor()) train_loader = DataLoader(train_dataset, batch_size=batch_size) test_loader = DataLoader(test_dataset, batch_size=batch_size) model = Model() sgd = SGD(model.parameters(), lr=1e-1) cost = CrossEntropyLoss() epoch = 100 for _epoch in range(epoch): model.train() for idx, (train_x, train_label) in enumerate(train_loader): label_np = np.zeros((train_label.shape[0], 10)) sgd.zero_grad() predict_y = model(train_x.float()) loss = cost(predict_y, train_label.long()) if idx % 10 == 0: print('idx: {}, loss: {}'.format(idx, loss.sum().item())) loss.backward() sgd.step() correct = 0 _sum = 0 model.eval() for idx, (test_x, test_label) in enumerate(test_loader): predict_y = model(test_x.float()).detach() predict_ys = np.argmax(predict_y, axis=-1) label_np = test_label.numpy() _ = predict_ys == test_label correct += np.sum(_.numpy(), axis=-1) _sum += _.shape[0] print('accuracy: {:.2f}'.format(correct / _sum)) torch.save(model, 'models/mnist_{:.2f}.pkl'.format(correct / _sum))

这段代码是一个简单的使用PyTorch训练MNIST数据集的示例代码。它包括了模型的定义、数据加载、训练和测试的过程。

首先，它导入了所需的库和模块，包括了模型定义(`Model`)、数据集(`mnist`)、数据加载器(`DataLoader`)、损失函数(`CrossEntropyLoss`)、优化器(`SGD`)以及数据转换(`ToTensor`)。

然后，在`if __name__ == '__main__':`的条件下，定义了一些训练所需的参数，如批量大小(`batch_size`)、训练集和测试集的数据加载器(`train_loader`和`test_loader`)。接着，创建了一个模型实例(`model`)和优化器实例(`sgd`)，以及定义了损失函数(`cost`)和训练的轮数(`epoch`)。

接下来，通过一个循环迭代训练数据集的每个批次。在每个批次中，将模型设置为训练模式(`model.train()`)，将优化器的梯度置零(`sgd.zero_grad()`)，通过模型前向传播得到预测结果(`predict_y`)，计算损失(`loss`)，并进行反向传播和参数更新(`loss.backward()`和`sgd.step()`)。同时，每训练10个批次，打印出当前的损失值。

接着，通过另一个循环对测试集进行预测，并计算准确率。在每个测试批次中，将模型设置为评估模式(`model.eval()`)，通过模型前向传播得到预测结果(`predict_y`)，将预测结果转换为类别标签(`predict_ys`)，并与真实标签进行比较，统计正确预测的数量(`correct`)和总样本数量(`_sum`)。最后，计算并打印出准确率。

最后，将训练好的模型保存到文件中，文件名中包含了准确率。

这段代码的作用是训练一个简单的模型来分类MNIST手写数字数据集，并保存训练好的模型。