import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets,transforms
import matplotlib.pyplot as plt

'''构建网络层'''
class ANN(nn.Module):
    def __init__(self):
        super(ANN,self).__init__() # 对继承自父类的属性进行初始化
        self.linear_1 = nn.Linear(in_features=28*28,out_features=512,bias=True)
        self.reLU_1 = nn.ReLU(inplace=False) # 如果设为True，会把输出直接覆盖到输入中，这样可以节省内存/显存
        self.linear_2 = nn.Linear(in_features=512,out_features=512,bias=True)
        self.reLU_2 = nn.ReLU(inplace=False)
        self.linear_3 = nn.Linear(in_features=512,out_features=256,bias=True)
        self.reLU_3 = nn.ReLU(inplace=False)
        self.linear_4 = nn.Linear(in_features=256,out_features=10,bias=True)

    def forward(self,x_para_1):
        x_reshape = torch.reshape(x_para_1,shape=(-1,28*28*1))
        x_linear_1 = self.linear_1(x_reshape) # 使用了python的__call__方法,而在__call__方法中调用了forward函数
        x_reLU_1 = self.reLU_1(x_linear_1)
        x_linear_2 = self.linear_2(x_reLU_1)
        x_reLU_2 = self.reLU_2(x_linear_2)
        x_linear_3 = self.linear_3(x_reLU_2)
        x_reLU_3 = self.reLU_3(x_linear_3)
        x_linear_4 = self.linear_4(x_reLU_3)
        return x_linear_4

'''主函数：训练模型'''
def main():
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    epochs = 10  # 迭代次数
    batch_size = 512  # 批量大小

    net = ANN().to(device)  # 将对象移动到目标设备
    criterion = nn.CrossEntropyLoss(reduce = None, weight = None, size_average = None, ignore_index = -100)  # 选择损失函数
    optimizer = optim.Adam(net.parameters(), weight_decay=0, amsgrad=False, lr=0.001, betas=(0.9, 0.999),
                           eps=1e-10)  # 选择优化方法

    transform = transforms.Compose([
        transforms.Resize(28),  # 设置输出图像大小
        transforms.ToTensor(),  # 数据转化为tensor
        transforms.Normalize((0.5,), (0.5,))  # 数据标准化
    ])
    dataset = datasets.MNIST("datasets/", train=True, download=False, transform=transform)  # 加载MNIST数据集
    dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True)  # 将整个数据集分成多个批次

    testdataset = datasets.MNIST("datasets/", train=False, download=False, transform=transform)
    testdataloader = torch.utils.data.DataLoader(testdataset, batch_size=batch_size, shuffle=True)

    losses = []
    for i in range(epochs):
        net.train()  # 打开滑动指数平均按钮，将batch上的mean和var近似成整个样本空间上的mean和var，在mini-batch训练集中使用
        print('epochs: %d' % i)
        for j, (input, target) in enumerate(dataloader):
            input, target = input.to(device), target.to(device)
            output = net(input)
            loss = criterion(output, target)

            optimizer.zero_grad()  # 将每一轮的梯度设为零
            loss.backward()  # 计算梯度，误差反向传播
            optimizer.step()  # 更新参数
            if j % 10 == 0:
                losses.append(loss.float())
                print("[epochs - %d - %d/%d]loss: %f" % (i, j, len(dataloader), loss.float()))
                '''实时绘图'''
                plt.clf()  # 清除当前 figure 上的内容
                plt.plot(losses)
                plt.savefig('loss.jpg')  # 保存图片
                plt.pause(0.01)  # 相当于plt.show()，但是只显示0.01秒

        with torch.no_grad():  # 测试集中的数据不需要计算梯度，也不会进行反向传播
            net.eval()  # 相对于net.train()，在测试集中使用
            correct = 0.
            total = 0.
            for input, target in testdataloader:
                input, target = input.to(device), target.to(device)
                output = net(input)
                _, predicted = torch.max(output.data, 1)  # 返回每一行中最大值的索引
                total += target.size(0)
                correct += (predicted == target).sum()
                accuracy = correct.float() / total
            print("[epochs - %d]Accuracy: %f" % (i + 1, (100 * accuracy)))
        torch.save(net, "models/net.pth") # 保存模型


if __name__ == "__main__":
    main()

loss.jpg

image.png