import timeimport torchfrom torch import nn, optimfrom torch.utils.data import DataLoaderfrom torchvision import datasets, transforms"""版本: python==3.7 pytorch==1.6"""# 超参LR = 0.001 ## 学习效率NUM_EPOCH = 2 ## 训练次数BATCH_SIZE = 50 ## 每批次训练 50 条数据device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') ## 使用 GPU 加速## 数据集准备train_dataset = datasets.MNIST( root='./dataset', ## 数据储存位置 train=True, ## True 为训练集 transform=transforms.ToTensor(), download=False ## 是否下载 MNIST 数据集(我已经下好了),没下载的写 True)train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)test_dataset = datasets.MNIST( root='./dataset', train=False, ## False为测试集 transform=transforms.ToTensor(), download=False)test_data = torch.unsqueeze(test_dataset.data, dim=1).type(torch.FloatTensor)[:8000].to(device) / 255. ## 测试集数据与训练集数据维度保持一致test_label = test_dataset.targets.to(device)[:8000]val_data = torch.unsqueeze(test_dataset.data, dim=1).type(torch.FloatTensor)[8000:].to(device) / 255. val_label = test_dataset.targets.to(device)[8000:]## 搭建 CNN 神经网络class CNN(nn.Module): def __init__(self): super(CNN, self).__init__() self.conv1 = nn.Sequential( nn.Conv2d( ## 卷积层 1 in_channels=1, ## 输入图片维度为 1 * 28 * 28 out_channels=16, ## 输出图片维度为 16 * 28 * 28 kernel_size=5, stride=1, padding=2 ), nn.ReLU(), nn.MaxPool2d(kernel_size=2) ## 池化层 经池化操作后图片维度变为 16 * 14 * 14 ) self.conv2 = nn.Sequential( nn.Conv2d( in_channels=16, out_channels=32, # 32 * 14 * 14 kernel_size=5, stride=1, padding=2 ), nn.ReLU(), nn.MaxPool2d(kernel_size=2)# 32 * 7 * 7 ) self.out = nn.Linear(32*7*7, 10) ## 全连接层 def forward(self, x): x = self.conv1(x) x = self.conv2(x) x = x.view(x.size()[0], -1) out = self.out(x) return outcnn = CNN().to(device) ## 将模型移植到 GPU 上loss_func = nn.CrossEntropyLoss() ## 损失函数optimizer = optim.Adam(cnn.parameters(), lr=LR) ## 优化器 ## 输出信息格式def print_info(num, str, line=True, start=True, end=True): if line: str = "-"*5 + " " + str + " " + "-"*5 blank = " " * ((num - len(str)) // 2) if start: print("\n" + "*" * num) print("\n%s%s%s\n" % (blank, str, blank)) if end: print("*" * num)## 开始训练start_0 = time.time()num = 69start_str = "Start training"print_info(num, start_str, end=False)for epoch in range(NUM_EPOCH): for step, (data, label) in enumerate(train_loader): data = data.to(device) ## 将数据移植到 GPU上 label = label.to(device) ## 将数据移植到 GPU上 out = cnn(data) loss = loss_func(out, label) optimizer.zero_grad() loss.backward() optimizer.step() if step % 100 == 0: with torch.no_grad(): test_out = cnn(test_data) pred_y = torch.max(test_out, 1)[1] accurancy = torch.sum(pred_y==test_label) / float(test_label.size()[0]) print("Epoch: [%s/%s] |, step: [%s/%s] |, loss: %.4f |, accuracy: %.2f%%" % (epoch+1, NUM_EPOCH, step, train_loader.__len__(), loss.data.item(), 100*accurancy.data.item()))### 烦死了,懒得写注释了, 不知道的去 scdn 或者 Google 或者 百度,不行再找我time_str = "Training time: %.2fs" % (time.time() - start_0)print_info(num, time_str, line=False, start=False, end=False)end_str = "The end of training"print_info(num, end_str, start=False)## 验证with torch.no_grad(): val_out = cnn(val_data) pred_y = torch.max(val_out, 1)[1] accuracy = torch.sum(pred_y == val_label) / float(val_label.size()[0]) start_str = "Start verification" print_info(num, start_str, start=False, end=False) val_str = "Validating set accurancy: %.2f%%" % (100*accurancy.data.item()) print_info(num, val_str, line=False, start=False, end=False) end_str = "The end of the validation" print_info(num, end_str, start=False)## savestart_str = "Start saving the model"print_info(num, start_str, start=False, end=False)start1 = time.time()torch.save(cnn, './models/cnn.pkl')save_str = "Save the elapsed time of model completely: %.2fs" % (time.time() - start1)print_info(num, save_str, line=False, start=False, end=False)start2 = time.time()torch.save(cnn.state_dict(), './models/cnn_params.pkl')save_str = "The time it takes to save model parameters: %.2fs" % (time.time() - start2)print_info(num, save_str, line=False, start=False, end=False)end_str = "The end of saving the model"print_info(num, end_str, start=False)
输出:
懒得写了,就这样。。。。。。