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)

输出：

懒得写了，就这样。。。。。。