Pytorch 5.3 自定义层、读写文件

自定义层

早在Pytorch 5.1 中我们就有讲过自定义层的实现和原理，具体可以参考我的前两篇文章。 这篇文章主要是补充带参数的层的自定义，和不带参数的层的自定义。

不带参数自定义层

class CenterLayer(nn.Module):  #  我们定义其为 a-a的均值 
    def __init__(self): 
        super().__init__() 
        
    def forward(self,X): 
        return X-X.mean()
    
layer = CenterLayer()(torch.tensor([1., 2., 3., 4., 5.])) 
layer 

out[1] : tensor([-2., -1.,  0.,  1.,  2.])

将层作为组件合并到更复杂的模型中

net = nn.Sequential(nn.Linear(16,256),CenterLayer()) 

Y = net(torch.rand(4,16)) 
Y.mean()

out[2] : tensor(-3.7253e-09, grad_fn=<MeanBackward0>)

带参数的层 1️⃣一个权重参数，2️⃣一个偏执参数，3️⃣使用线性修正单元为激活函数

class MyLinear(nn.Module): 
    def __init__(self,in_units,out_units): 
        super().__init__() 
        self.weight = nn.Parameter(data=torch.randn(in_units,out_units)) 
        self.bias = nn.Parameter(data=torch.randn(out_units,)) 
        
    def forward(self,X): 
        linear = torch.matmul(X,self.weight.data) + self.bias.data 
        return F.relu(linear) 
    
linear = MyLinear(in_units=5,out_units=3) 
linear(torch.rand(2,5))

linear.weight
out[3]: Parameter containing:
tensor([[ 1.0556,  0.9521, -1.4115],
        [-1.2977, -0.2838, -0.1454],
        [-0.1564,  0.2859,  0.2608],
        [ 1.4347, -0.8297,  0.4063],
        [ 0.0141,  0.0456, -0.1127]], requires_grad=True)

读写文件

保存文件函数：

torch.save(obj,f: Union[str, os.PathLike, BinaryIO, IO[bytes]],...) # obj:即将保存的参数， f:保存的文件名/地址 

加载文件函数：

torch.load(f, map_location=None, pickle_module=pickle, **pickle_load_args) # f:加载的文件名，地址

保存和读取张量

x = torch.arange(4)
torch.save(x,'x-file') # ------> 发现根目录下多了一个 x-file 文件
x2 = torch.load('x-file') 
x2 

tensor([0, 1, 2, 3])

保存一个张量列表 (为节省空间不打印输出结果)

y = torch.zeros(4)
torch.save([x, y],'x-files')
x2, y2 = torch.load('x-files')

写入或读取从字符串映射到张量的字典

mydict = {'x':x,'y':y}
torch.save(mydict,'mydict') 
mydict2 = torch.load('mydict')  
mydict2

{'x': tensor([0, 1, 2, 3]), 'y': tensor([0., 0., 0., 0.])}

保存模型参数

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)

torch.save(net.state_dict(), 'mlp.params')

加载模型参数

clone = MLP()
clone.load_state_dict(torch.load('mlp.params'))  # <All keys matched successfully> 
clone.eval() # 评估模式

MLP(
  (hidden): Linear(in_features=20, out_features=256, bias=True)
  (output): Linear(in_features=256, out_features=10, bias=True)
)

# 我们参数一样，传入的数据一样，那么计算出来的结果应该是一样的 
Y_clone = clone(X) 
Y_clone == Y 
'''Out: 
tensor([[True, True, True, True, True, True, True, True, True, True],
        [True, True, True, True, True, True, True, True, True, True]])
'''