import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

下载训练集

train_dataset = datasets.MNIST(root='./',

                           train=True,                           transform=transforms.ToTensor(),                           download=True)

下载测试集

test_dataset = datasets.MNIST(root='./',

                          train=False,                          transform=transforms.ToTensor(),                          download=True)

批次大小

batch_size = 64

装载训练集

train_loader = DataLoader(dataset=train_dataset,

                      batch_size=batch_size,                      shuffle=True)

装载测试集

test_loader = DataLoader(dataset=test_dataset,

                     batch_size=batch_size,                     shuffle=True)

for i, data in enumerate(train_loader):

# 取得数据和对应的标签inputs, labels = dataprint(inputs.shape)print(labels.shape)break

定义网络结构

class Net(nn.Module):

def __init__(self):    super(Net, self).__init__()    # 卷积层1    # Conv2d 参数1：[金属期货](https://www.gendan5.com/cf/mf.html)输出通道数，黑白图片为1，黑白为3 参数2：输入通道数，生成32个特色图 参数3:5*5卷积窗口 参数4：步长1 参数5：padding补2圈0（3*3卷积窗口填充1圈0,5*5填充2圈0）    # 应用ReLU激活函数 池化窗口大小2*2，步长2    self.conv1 = nn.Sequential(nn.Conv2d(1, 32, 5, 1, 2), nn.ReLU(), nn.MaxPool2d(2, 2))    # 卷积层2 输出32个特色图 输入64个特色图    self.conv2 = nn.Sequential(nn.Conv2d(32, 64, 5, 1, 2), nn.ReLU(), nn.MaxPool2d(2, 2))    # 全连贯层1 输出64*7*7（原先为28，每次池化/2），1000    self.fc1 = nn.Sequential(nn.Linear(64 * 7 * 7, 1000), nn.Dropout(p=0.4), nn.ReLU())    # 全连贯层2 输入10个分类，并转化为概率    self.fc2 = nn.Sequential(nn.Linear(1000, 10), nn.Softmax(dim=1))def forward(self, x):    # 卷积层应用4维的数据    # 批次数量64 黑白1 图片大小28*28    # ([64, 1, 28, 28])    x = self.conv1(x)    x = self.conv2(x)    # 全连贯层对2维数据进行计算    x = x.view(x.size()[0], -1)    x = self.fc1(x)    x = self.fc2(x)    return x

LR = 0.0003

定义模型

model = Net()

定义代价函数

entropy_loss = nn.CrossEntropyLoss()

定义优化器

optimizer = optim.Adam(model.parameters(), LR)
def train():

model.train()for i, data in enumerate(train_loader):    # 取得数据和对应的标签    inputs, labels = data    # 取得模型预测后果，（64，10）    out = model(inputs)    # 穿插熵代价函数out(batch,C),labels(batch)    loss = entropy_loss(out, labels)    # 梯度清0    optimizer.zero_grad()    # 计算梯度    loss.backward()    # 修改权值    optimizer.step()

def test():

model.eval()correct = 0for i, data in enumerate(test_loader):    # 取得数据和对应的标签    inputs, labels = data    # 取得模型预测后果    out = model(inputs)    # 取得最大值，以及最大值所在的地位    _, predicted = torch.max(out, 1)    # 预测正确的数量    correct += (predicted == labels).sum()print("Test acc: {0}".format(correct.item() / len(test_dataset)))correct = 0for i, data in enumerate(train_loader):    # 取得数据和对应的标签    inputs, labels = data    # 取得模型预测后果    out = model(inputs)    # 取得最大值，以及最大值所在的地位    _, predicted = torch.max(out, 1)    # 预测正确的数量    correct += (predicted == labels).sum()print("Train acc: {0}".format(correct.item() / len(train_dataset)))

for epoch in range(0, 10):

print('epoch:', epoch)train()test()