LeNet
早期成功的神经网络。
先使用卷积层来学习图片空间信息。
然后使用全连接层转换到类别空间。

#In[]
'''
LeNet,上世纪80年代的产物,最初为了手写识别设计
'''
from d2l import torch as d2l
import torch 
from torch import nn
from torch.nn.modules.loss import CrossEntropyLoss

from torch.utils import data
import torchvision
from torchvision import transforms
import matplotlib.pyplot as plt
import Common_functions


'''
LeNet:
两个卷积层,两个池化层,三个线性层
假定为MNIST设计,输入为(batch_size,1,28,28)
'''

class Reshape(torch.nn.Module):
    def forward(self,x):
        return x.view(-1,1,28,28)

net = nn.Sequential(
    nn.Conv2d(in_channels=1,out_channels=6,kernel_size=(5,5),padding=2),nn.Sigmoid(), #输出:(6,28,28)
    nn.AvgPool2d(kernel_size=(2,2)), #不指定stride默认不重叠 输出(6,14,14)
    nn.Conv2d(6,16,kernel_size=(5,5)),nn.Sigmoid(),#输出(16,10,10)
    nn.AvgPool2d(kernel_size=(2,2)),#输出(16,5,5)
    nn.Flatten(),
    nn.Linear(16*5*5,120),nn.Sigmoid(),#
    nn.Linear(120,84),nn.Sigmoid(),
    nn.Linear(84,10)
)


X=torch.rand(size=(1,1,28,28),dtype=torch.float32)
for layer in net:
    X=layer(X)
    print(layer.__class__.__name__,'output shape: \t',X.shape)

#In[]


batch_size = 256
train_iter,test_iter=d2l.load_data_fashion_mnist(batch_size=batch_size)



#对evaluate_accuracy函数进行轻微修改
#使用GPU计算模型在数据集上的精度
#计算网络在测试数据集上面的准确率
#由于完整的测试数据集位于内存中，因此在模型使用GPU预测测试数据集之前，我们需要将其复制到显存中。
def evaluate_accuracy_gpu(net,data_iter,device=None):
    if isinstance(net,nn.Module):
        net.eval() #网络用于测试数据
        if not device:
            device = next(iter(net.parameters())).device #如果没有指定device设备，device设备则使用第一层网络参数的设备
    accumulator = d2l.Accumulator(2) #累加器里面包含两个元素
    for X,y in data_iter:
        if isinstance(X,list):
            X = [x.to(device) for x in X] #X为list类型时，需要加X里面每个元素都复制到device设备上面来
        else:
            X = X.to(device)
        y = y.to(device)
        accumulator.add(d2l.accuracy(net(X),y),y.numel()) #累加器第一个元素为在每一个batch_size中预测准确的个数，第二个元素为每一个batch_size中样本总数目，然后依次循环累加，得到测试数据集上面预测准确的总数目，以及数据集总数目
    return accumulator[0]/accumulator[1] #算出模型预测准确率


def train_ch6(net,train_iter,test_iter,num_epochs,lr,device):
    def init_weights(m):#手动初始化模型参数
        if type(m) == nn.Linear or type(m) == nn.Conv2d:
            nn.init.xavier_uniform_(m.weight) #使用xavier_uniform分布初始化参数
    net.apply(init_weights)
    net.to(device)#将模型复制到gpu上面
    print('training on',device)
    loss = nn.CrossEntropyLoss() #定义loss
    optim = torch.optim.SGD(net.parameters(),lr=lr) #定义优化器
    animator = d2l.Animator(xlabel='epoch',xlim=[1,num_epochs],legend=['train_loss','train_acc','test_acc'])
    timer = d2l.Timer()
    num_batches = len(train_iter)
    for epoch in range(num_epochs):
        net.train()#模型开始训练，需要放在第一层循环里面，因为后面evaluate_accuracy_gpu()函数里面有net.eval()，将模型改变为测试状态，因此需要在每一个循环epoch后面手动再加上模型开始处于训练状态
        accumulator = d2l.Accumulator(3) #累加器
        for i,(X,y) in enumerate(train_iter):
           timer.start()
           optim.zero_grad()
           X = X.to(device)#将X复制到gpu上面
           y = y.to(device) #将y复制到gpu上面
           y_hat = net(X) #得到模型训练后的输出标签y_hat
           l = loss(y_hat,y)#计算每一个batch_size的loss
           l.backward() #计算梯度
           optim.step() #使用优化器更新模型参数
           with torch.no_grad():#不需要模型梯度
               accumulator.add(l*X.shape[0],d2l.accuracy(y_hat,y),X.shape[0])
           timer.stop()
           train_loss = accumulator[0]/accumulator[2] #从累加器里面获得所有训练集的loss之和
           train_acc = accumulator[1]/accumulator[2] #从累加器里面获得所有训练集的准确数之和
           if (i+1) % (num_batches // 5) == 0 or i == num_batches-1:
               animator.add(epoch+(i+1)/num_batches,(train_loss,train_acc,None))
        test_accuracy = evaluate_accuracy_gpu(net,test_iter) #每次训练完一个epoch后的模型用于测试数据集上面计算测试精确度
        animator.add(epoch+1,(None,None,test_accuracy))
    print(f'模型训练完最后一轮时 train_loss:{train_loss},train_acc:{train_acc},test_acc:{test_accuracy}')
    print(f'{num_epochs*accumulator[2]/timer.sum()}examples/second on {str(device)}')#打印出模型每秒能处理多少个样本数

lr,num_epochs= 0.9,10
train_ch6(net,train_iter=train_iter,test_iter=test_iter,lr=lr,num_epochs=num_epochs,device=d2l.try_gpu())
'''
输出结果：
模型训练完最后一轮时 train_loss:0.4322478462855021,train_acc:0.8396666666666667,test_acc:0.8163
55954.65804440994examples/second on cuda:0
'''








#训练
if torch.cuda.is_available():
    device = "cuda:0"
else:
    device = "cpu"
device = torch.device(device)

Common_functions.train_device(net,train_iter,test_iter,lr=0.9,device=device)
# %%

plt.show()