##神经网络结构为 https://pytorch.org/tutorials/_images/mnist.png

import torch
import torch.nn as nn  # nn -- neural network， 使用它创建卷积层、全连接层
import torch.nn.functional as F  # nn 提供的函数库， 使用函数方式创建池化层和激活函数

## 构建模型类
class Net(nn.Module):

    def __init__(self):
        super(Net, self).__init__()
        # 输入图像channel：1；输出channel：6；5x5卷积核
        self.conv1 = nn.Conv2d(1, 6, 5)  # 卷积层1 操作：1 * 32 * 32 -- 5x5 -> 6 * 28 * 28
        self.conv2 = nn.Conv2d(6, 16, 5)  # 卷积层2 操作：6 * 14 * 14 -- 5x5 -> 16 * 10 * 10
        # an affine operation: y = Wx + b 
        self.fc1 = nn.Linear(16 * 5 * 5, 120)  # 全连接层1 16 channel * 5 height * 5 width --> 120 * 1
        self.fc2 = nn.Linear(120, 84)  # 全连接层2 120 --> 84
        self.fc3 = nn.Linear(84, 10)  # 全连接层3 84 --> 10

    # 正向传播的过程：
    # conv1 --> relu --> max_poo12d --> conv2 --> relu --> max_pool2d --> view
    # --> fc1 --> relu --> fc2 --> relu --> fc3 --> out 
    def forward(self, x):
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))  # 池化层用于对数据采样，（2，2）--> 1
        # 如果是方阵,则可以只使用一个数字进行定义
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)  # relu() 函数用于增加神经网络的非线性，在每一层后使用
        x = x.view(-1, self.num_flat_features(x))  # view() 函数用于改变数据维度，这里得到一个 16 * 25 结构的数据
        x = F.relu(self.fc1(x))  # 全连接层
        x = F.relu(self.fc2(x))  # 全连接层
        x = self.fc3(x)  # # 全连接层，最后一层不经过 relu
        return x

    def num_flat_features(self, x):  # 返回一个 channel 中的 width * height
        size = x.size()[1:]  # 除去批处理维度的其他所有维度
        num_features = 1
        for s in size:
            num_features *= s
        return num_features

## 模型实例化
net = Net()  # 实例化一个Net对象
print(net)  # 打印对象net的结构
# 查看模型可学习参数
params = list(net.parameters())  # net.parameters() 返回模型的可学习参数
print(len(params))  # 输出 10，即每一层有两个参数 -- weight、bias
print(params[0].shape)  # conv1 的 weight 参数的结构为 torch.Size([6, 1, 5, 5])

## 进行正向传播
input = torch.randn(1, 1, 32, 32)  # 构造随机输入数据 
output = net(input)  # 正向传播
target = torch.randn(10)  # 构造随机目标数据
target = target.view(1, -1)  # 目标函数升维到 2 维
criterion = nn.MSELoss()  # 构造损失函数
loss = criterion(output, target)  # 计算模型的 loss 值

## 进行梯度计算和参数更新（手动更新参数）
net.zero_grad()  # 清零所有参数的梯度缓存 -- 否则 pytorch 计算的梯度将会累乘到上次的缓存上
# 谁调用 backward()，更新的梯度值就指向令谁下降的方向
loss.backward()  # 进行随机梯度的反向传播，从 loss 开始向可学习参数计算累计梯度，最后的梯度下降朝着 loss 下降的方向进行
learning_rate = 0.01  # learning_rate 用于控制梯度下降速度
for f in net.parameters():
    f.data.sub_(f.grad.data * learning_rate)  # 更新可学习参数的参数值

## 使用优化器进行参数更新
import torch.optim as optim
# 创建优化器 (optimizer)，可以使用优化器更好的更新可学习参数的值
optimizer = optim.SGD(net.parameters(), lr=0.01)  # 常见的优化器有SGD、Nesterov-SGD、Adam、RMSProp等，传入可学习参数和学习速度
optimizer.zero_grad()  # 清零梯度缓存
output = net(input)  # 正向传播
loss = criterion(output, target)  # 损失值计算
loss.backward()  # loss 调用 backward(), 梯度指向使loss下降的方向
optimizer.step()  # 更新参数