实验一程序.zip

import torch from torch import nn from torch import optim import torch.nn.functional as F import torchvision train_batch_size = 128 test_batch_size = 10000 img_size = 28 def get_dataloader(train=True): assert isinstance(train, bool), "train 必须是bool类型" # 准备数据集，其中0.1307，0.3081为MNIST数据的均值和标准差，这样操作能够对其进行标准化 # 因为MNIST只有一个通道（黑白图片）,所以元组中只有一个值 dataset = torchvision.datasets.MNIST('/data', train=train, download=True, transform=torchvision.transforms.Compose([ torchvision.transforms.ToTensor(), torchvision.transforms.Normalize((0.1307,), (0.3081,)), ])) # 准备数据迭代器 batch_size = train_batch_size if train else test_batch_size dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True) return dataloader class MnistNet(nn.Module): def __init__(self): super(MnistNet, self).__init__() self.fc1 = nn.Linear(28 * 28 * 1, 28) self.fc2 = nn.Linear(28, 10) def forward(self, x): x = x.view(-1, 28 * 28 * 1) x = self.fc1(x) # [batch_size,28] x = F.relu(x) # [batch_size,28] x = self.fc2(x) # [batch_size,10] return x # return F.log_softmax(x, dim=-1) mnist_net = MnistNet() optimizer = optim.Adam(mnist_net.parameters(), lr=0.001) train_loss_list = [] train_count_list = [] def train(epoch): mode = True mnist_net.train(mode=mode) train_dataloader = get_dataloader(train=mode) for idx, (data, target) in enumerate(train_dataloader): optimizer.zero_grad() output = mnist_net(data) loss = F.cross_entropy(output, target) # 对数似然损失 loss.backward() optimizer.step() if idx % 10 == 0: print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format( epoch, idx * len(data), len(train_dataloader.dataset), 100. * idx / len(train_dataloader), loss.item())) train_loss_list.append(loss.item()) train_count_list.append(idx * train_batch_size + (epoch - 1) * len(train_dataloader)) torch.save(mnist_net.state_dict(), "./mnist_net.pkl") torch.save(optimizer.state_dict(), './mnist_optimizer.pkl') def test1(): test_loss = 0 correct = 0 mnist_net.eval() test_dataloader = get_dataloader(train=False) with torch.no_grad(): for data, target in test_dataloader: output = mnist_net(data) output_=F.softmax(output) test_loss += F.cross_entropy(output, target, reduction='sum').item() pred = output.data.max(1, keepdim=True)[1] # 获取最大值的位置,[batch_size,1] correct += pred.eq(target.data.view_as(pred)).sum() test_loss /= len(test_dataloader.dataset) print('\nTest set: Avg. loss: {:.4f}, Accuracy: {}/{} ({:.2f}%)\n'.format( test_loss, correct, len(test_dataloader.dataset), 100. * correct / len(test_dataloader.dataset))) return pred,target,output_ if __name__ == '__main__': test1() for i in range(1): # 模型训练5轮 train(i) y_pre,y_ture,scores=test1() y_pre=torch.squeeze(y_pre) torch.save(y_pre,'./data/y_pre.pt') torch.save(y_ture,'./data/y_true.pt') torch.save(scores,'./data/scores.pt')