import torch
from torch import nn
from torch import optim
import torch.nn.functional as F

# Linear 继承 nn.Module
class Linear(nn.Module):
    def __init__(self,in_dim,out_dim):
        super(Linear,self).__init__()  # 调用nn.Module构造函数
        self.w = nn.Parameter(torch.randn(in_dim,out_dim))  #nn.Parameter 构造学习参数,初始参数随机
        self.b = nn.Parameter(torch.randn(out_dim))
    def forward(self,x):     #前向传播
        x = x.matmul(self.w)
        y = x+self.b.expand_as(x)
        return y
# 感知机，继承nn.Module 调用Linear
class Perception(nn.Module):
    def __init__(self,in_dim,hid_dim,out_dim):
        super(Perception,self).__init__()
        self.layer1 = Linear(in_dim,hid_dim)
        self.layer2 = Linear(hid_dim,out_dim)
    def forward(self,x):     #前向传播
        x = self.layer1(x)
        y = torch.sigmoid(x) #激活函数
        y = self.layer2(y)
        y = torch.sigmoid(y) #激活函数
        return y

    

# 实例化一个网络，实现2分类
perception = Perception(2,3,2)
perception

# .named_parameters() 可以返回学习参数的迭代器，包含参数名、参数值
for name ,parameter in perception.named_parameters():
    print(name,‘-----‘,parameter)

# 生成输入数据：4个2维 数据
data = torch.randn(4,2)
print(data)

output = perception(data)  #直接会调用forward
print(output)

# 对于需要学习参数的层最好使用nn.Module
# nn.functional 定义的网络层不可自动学习参数
# nn.Sequential() 模块，提供快速构建简单的网络
class Perceptionv2(nn.Module):
    def __init__(self,in_dim,hid_dim,out_dim):
        super(Perceptionv2,self).__init__()
        self.layer=nn.Sequential(nn.Linear(in_dim,hid_dim),nn.Sigmoid(),nn.Linear(hid_dim,out_dim),nn.Sigmoid())
    def forward(self,x):
        y=self.layer(x)
        return y

# 定义一个Perceptionv2 对象，100维数据，输出维度为10
model = Perceptionv2(100,1000,10).cuda()
model

# 初始化输入数据，100维
input = torch.randn(100).cuda()
output2 = model(input)
print(output2.shape)

# 损失函数
# functional  和 torch.nn 中都包含有损失函数，后者多为对象
# 因为损失函数通常不含有可学习的参数，因此两者都可以使用
import torch.nn.functional as F

print(output) #4个体样本二分类
# 获取标签 4个样本
label = torch.Tensor([1,0,1,1]).long()

# nn 中的交叉熵损失类
criterion = nn.CrossEntropyLoss()
loss_nn = criterion(output,label)
print(loss_nn)

# torch.nn.functiona 中的cross_entropy 交叉熵损失函数
loss_functional = F.cross_entropy(output,label)
print(loss_functional)

# 优化器 nn.optim
from torch import optim
#如下方式调用
#optimizer = optim.SGD(model.parameters(),lr=0.001,momentum=0.9)
#optimizer = optim.Adam([var1,var2],lr=0.0001)

# 使用一个三层的感知机，进行示例
class MLP(nn.Module):
    def __init__(self,in_dim,hid_dim1,hid_dim2,out_dim):
        super(MLP,self).__init__()
        # 通过Sequential 快速搭建
        self.layer = nn.Sequential(nn.Linear(in_dim,hid_dim1),nn.ReLU(),nn.Linear(hid_dim1,hid_dim2),nn.ReLU(),nn.Linear(hid_dim2,out_dim),nn.ReLU())
    def forward(self,x):
        x = self.layer(x)
        return x

# 实例化一个模型
model = MLP(28*28,300,200,10)
print(model)
print(model.parameters()) #返回一个参数迭代器

# 采用SGD优化器,学习率为0.01
optimizer = optim.SGD(params=model.parameters(),lr=0.01) #和模型参数绑定，制定学习率等相关参数
# 输入数据，10个 每个维度为28*28
data = torch.randn(10,28*28)
output = model(data)  #output
print(output.shape)
# 10个样本 10分类
label = torch.Tensor([1,0,4,7,9,3,4,5,3,2]).long()
# 求损失
criterion = nn.CrossEntropyLoss()

loss = criterion(output,label)
print(loss)
print(type(loss))
# 在每次优化之前都需要，清空梯度
optimizer.zero_grad()
# 损失函数的反向传播
loss.backward()
# 梯度更新，利用优化器
optimizer.step()
# 上面的代码中使用 optim.SGD，并为所用层的参数，指定了统一的学习率
# 也可以为特殊的参数层，单独指定学习率
# 也可以在迭代次数超过一定的数量后，重新赋予optim优化器新的学习率