# 回顾一下多层感知机
import torch
from torch import nn
from torch.nn import functional as F
net = nn.Sequential(nn.Linear(20,256),nn.ReLU(),nn.Linear(256,10))
X = torch.rand(2,20)
net(X)

class MLP(nn.Module):
    def __init__(self):
        super().__init__()  # 调用父类的__init__函数
        self.hidden = nn.Linear(20,256)
        self.out = nn.Linear(256,10)
        
    def forward(self, X):
        return self.out(F.relu(self.hidden(X)))
    
# 实例化多层感知机的层，然后在每次调用正向传播函数调用这些层
net = MLP()
X = torch.rand(2,20)
net(X)

class MySequential(nn.Module):
    def __init__(self, *args):
        super().__init__()
        for block in args:
            self._modules[block] = block # block 本身作为它的key，存在_modules里面的为层，以字典的形式
            
    def forward(self, X):
        for block in self._modules.values():
            print(block)
            X = block(X)
        return X
    
net = MySequential(nn.Linear(20,256),nn.ReLU(),nn.Linear(256,10))
X = torch.rand(2,20)
net(X)

# 在正向传播函数中执行代码
class FixedHiddenMLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.rand_weight = torch.rand((20,20),requires_grad=False)
        self.linear = nn.Linear(20,20)
    
    def forward(self, X):
        X = self.linear(X)
        X = F.relu(torch.mm(X, self.rand_weight + 1))
        X = self.linear(X)
        while X.abs().sum() > 1:
            X /= 2
        return X.sum()
    
net = FixedHiddenMLP()
X = torch.rand(2,20)
net(X)

# 混合代培各种组合块的方法
class NestMLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(20,64),nn.ReLU(),
                                nn.Linear(64,32),nn.ReLU())
        self.linear = nn.Linear(32,16)
        
    def forward(self, X):
        return self.linear(self.net(X))
    
chimear = nn.Sequential(NestMLP(),nn.Linear(16,20),FixedHiddenMLP())
X = torch.rand(2,20)
chimear(X)

# 首先关注具有单隐藏层的多层感知机
import torch
from torch import nn

net = nn.Sequential(nn.Linear(4,8),nn.ReLU(),nn.Linear(8,1))
X = torch.rand(size=(2,4))
print(net(X))
print(net[2].state_dict()) # 访问参数，net[2]就是最后一个输出层
print(type(net[2].bias)) # 目标参数
print(net[2].bias)
print(net[2].bias.data)
print(net[2].weight.grad == None) # 还没进行反向计算，所以grad为None
print(*[(name, param.shape) for name, param in net[0].named_parameters()])  # 一次性访问所有参数         
print(*[(name, param.shape) for name, param in net.named_parameters()])  # 0是第一层名字，1是ReLU，它没有参数
print(net.state_dict()['2.bias'].data) # 通过名字获取参数

# 从嵌套块收集参数
def block1():
    return nn.Sequential(nn.Linear(4,8),nn.ReLU(),nn.Linear(8,4),nn.ReLU())

def block2():
    net = nn.Sequential()
    for i in range(4):
        net.add_module(f'block{i}',block1()) # f'block{i}' 可以传一个字符串名字过来，block2可以嵌套四个block1                                      
    return net

rgnet = nn.Sequential(block2(), nn.Linear(4,1))
print(rgnet(X))
print(rgnet)

net = nn.Sequential(nn.Linear(4,8),nn.ReLU(),nn.Linear(8,1))

def init_normal(m):
    if type(m) == nn.Linear:
        nn.init.normal_(m.weight, mean=0, std=0.01) # 下划线表示把m.weight的值替换掉   
        nn.init.zeros_(m.bias)
        
net.apply(init_normal) # 会递归调用 直到所有层都初始化
print(net[0].weight.data[0])
print(net[0].bias.data[0])

net = nn.Sequential(nn.Linear(4,8),nn.ReLU(),nn.Linear(8,1))

def init_constant(m):
    if type(m) == nn.Linear:
        nn.init.constant_(m.weight,1)
        nn.init.zeros_(m.bias)
        
net.apply(init_constant)
print(net[0].weight.data[0]) 
print(net[0].bias.data[0])

# 对某些块应用不同的初始化
def xavier(m):
    if type(m) == nn.Linear:
        nn.init.xavier_uniform_(m.weight)
        
def init_42(m):
    if type(m) == nn.Linear:
        nn.init.constant_(m.weight, 42)
        
net[0].apply(xavier)
net[2].apply(init_42)
print(net[0].weight.data[0])
print(net[2].weight.data)

# 自定义初始化
def my_init(m):
    if type(m) == nn.Linear:
        print("Init",*[(name, param.shape) for name, param in m.named_parameters()][0])  # 打印名字是啥，形状是啥       
        nn.init.uniform_(m.weight, -10, 10)
        m.weight.data *= m.weight.data.abs() >=  5 # 这里*=的代码相当于先计算一个布尔矩阵(先判断>=)，然后再用布尔矩阵的对应元素去乘以原始矩阵的每个元素。保留绝对值大于5的权重，不是的话就设为0

net.apply(my_init)
print(net[0].weight[:2])
net[0].weight.data[:] += 1 # 参数替换
net[0].weight.data[0,0] = 42
print(net[0].weight.data[0])

# 参数绑定
shared = nn.Linear(8,8)
net = nn.Sequential(nn.Linear(4,8),nn.ReLU(),shared,nn.ReLU(),shared,nn.ReLU(),nn.Linear(8,1))  # 第2个隐藏层和第3个隐藏层是share权重的，第一个和第四个是自己的  
net(X)
print(net[2].weight.data[0] == net[4].weight.data[0])
net[2].weight.data[0,0] = 100
print(net[2].weight.data[0] == net[4].weight.data[0])

# 构造一个没有任何参数的自定义层
import torch
import torch.nn.functional as F
from torch import nn

class CenteredLayer(nn.Module):
    def __init__(self):
        super().__init__()
        
    def forward(self, X):
        return X - X.mean()
    
layer = CenteredLayer()
print(layer(torch.FloatTensor([1,2,3,4,5])))

# 将层作为组件合并到构建更复杂的模型中
net = nn.Sequential(nn.Linear(8,128),CenteredLayer())
Y = net(torch.rand(4,8))
print(Y.mean())

# 带参数的图层
class MyLinear(nn.Module):
    def __init__(self, in_units, units):
        super().__init__()
        self.weight = nn.Parameter(torch.randn(in_units,units)) # nn.Parameter使得这些参数加上了梯度    
        self.bias = nn.Parameter(torch.randn(units,))

    def forward(self, X):
        linear = torch.matmul(X, self.weight.data) + self.bias.data
        return F.relu(linear)
    
dense = MyLinear(5,3)
print(dense.weight)

# 使用自定义层直接执行正向传播计算
print(dense(torch.rand(2,5)))
# 使用自定义层构建模型
net = nn.Sequential(MyLinear(64,8),MyLinear(8,1))
print(net(torch.rand(2,64)))

# 加载和保存张量
import torch
from torch import nn
from torch.nn import functional as F

x = torch.arange(4)
torch.save(x, 'x-file')
x2 = torch.load("x-file")
print(x2)

#存储一个张量列表，然后把它们读回内存
y = torch.zeros(4)
torch.save([x,y],'x-files')
x2, y2 = torch.load('x-files')
print(x2)
print(y2)

# 写入或读取从字符串映射到张量的字典
mydict = {'x':x,'y':y}
torch.save(mydict,'mydict')
mydict2 = torch.load('mydict')
print(mydict2)

# 加载和保存模型参数
class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.hidden = nn.Linear(20,256)
        self.output = nn.Linear(256,10)
    
    def forward(self, x):
        return self.output(F.relu(self.hidden(x)))
    
net = MLP()
X = torch.randn(size=(2,20))
Y = net(X)

# 将模型的参数存储为一个叫做"mlp.params"的文件
torch.save(net.state_dict(),'mlp.params')

# 实例化了原始多层感知机模型的一个备份。直接读取文件中存储的参数
clone = MLP() # 必须要先声明一下，才能导入参数
clone.load_state_dict(torch.load("mlp.params"))
print(clone.eval()) # eval()是进入测试模式

Y_clone = clone(X)
print(Y_clone == Y)