自定义层和读写文件

Psycho_MrZhang2025-10-20 11:53

自定义层

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

python 复制代码

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()
layer(torch.FloatTensor([1, 2, 3, 4, 5]))

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

python 复制代码

net = nn.Sequential(nn.Linear(8, 128), CenteredLayer())

Y = net(torch.rand(4m 8))
Y.mean()

带参数的层

python 复制代码

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.data
		return F.relu(linear)

dense = MyLinear(5, 3)
dense.weight

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

python 复制代码

dense(torch.rand(2, 5))

读写文件

python 复制代码

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')
x2 == x

存储一个张量列表

python 复制代码

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

写入或读取字典

python 复制代码

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

加载和保存模型参数

python 复制代码

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)

将模型存储为文件

python 复制代码

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

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