自定义层和读写文件

自定义层

自定义一个没有任何参数的层

import torch
import torch.nn.functional as F
from torch import nnclass CenteredLayer(nn.Module):def __init__(self):super().__init__()def forward(self, X):return X - X.mean()layer = CenteredLayer()
layer(torch.FloatTensor([1, 2, 3, 4, 5]))

将层作为组件和冰岛构建更复杂的模型中

net = nn.Sequential(nn.Linear(8, 128), CenteredLayer())Y = net(torch.rand(4m 8))
Y.mean()

带参数的层

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

使用自定义的层执行传播计算

dense(torch.rand(2, 5))

读写文件

import torch
from torch import nn
from torch.nn import functional as Fx = torch.arange(4)
torch.save(x, 'x-file')
x2 = torch.load('x-file')
x2 == x

存储一个张量列表

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')

加载和保存模型参数

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')# 保存参数后需要我们自己保存MLP的定义, 需要有定义才能加载
clone = MLP()
clone.load_state_dict(torch.load('mlp.params'))
clone.eval()