"""
演示
    搭建神经网络，这里是一个简单的全连接神经网络
这里的全连接神经网络组成：
    隐藏层1：nn.Linear(3,3),sigmoid,采用默认参数初始化
    隐藏层2：nn.Linear(3,2),relu,采用默认参数初始化
    输出层: nn.Linear(2,2),采用softmax做数据归一化

神经网络训练的一般工作流
    1.准备数据集
    2.搭建神经网络
        继承nn.Module父类
        定义__init__方法，定义网络层
        定义forward方法，实现前向传播
    3.模型训练
        前向传播
        计算损失
        梯度清零
        反向传播
        更新参数
    4.模型测试

"""
# 导包
import torch
from torch import nn
from torchsummary import summary  # 计算模型参数,查看模型结构, pip install torchsummary -i https://mirrors.aliyun.com/pypi/simple/

# 设置设备,优先级 CUDA > CPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"设备:{device}")

# 1.定义模型类，搭建神经网络
class MyModel(nn.Module):
    # 1.定义__init__方法，定义网络层
    def __init__(self):
        # 1.初始化父类成员
        super().__init__()
        # 2.定义网络层
        # 2.1 定义隐藏层1
        self.linear1 = nn.Linear(3,3)
        # 2.2 定义隐藏层2
        self.linear2 = nn.Linear(3,2)
        # 2.3 定义输出层
        self.linear3 = nn.Linear(2,2)
        # 3.参数初始化，一般不需要人为初始化，采用默认初始化即可

    # 2.定义forward方法，实现前向传播
    def forward(self,x):
        # x: 2D:(batch_size, 3)
        # 1.把数据输入隐藏层1
        # x = self.linear1(x)
        # # 经过sigmoid
        # x = torch.sigmoid(x)
        # 合并起来
        x = torch.sigmoid(self.linear1(x))
        # 2.把数据输入隐藏层2
        x = torch.relu(self.linear2(x))
        # 3.把数据输入输出层
        x = torch.softmax(self.linear3(x),dim=-1)   # (batch_size,2)
        # 4.返回结果
        return x

# 测试
if __name__ == '__main__':
    # 1.创建模型对象
    model = MyModel()
    print(model)
    # 迁移模型到设备，如果有CUDA设备，迁移到CUDA设备上，否则迁移到CPU设备上
    model = model.to(device)
    # 2.查看模型的参数数量
    summary(model, (3,))
    ...