PyTorch中定义可学习参数时的坑-Toy模板网

这篇具有很好参考价值的文章主要介绍了PyTorch中定义可学习参数时的坑。希望对大家有所帮助。如果存在错误或未考虑完全的地方，请大家不吝赐教，您也可以点击"举报违法"按钮提交疑问。

当需要在模型运行时定义可学习参数时（常见场景：参数的维度由每一层的维度定），我们就需要用这样的写法来实现：

class model(torch.nn.Module):
	def __init__(self):
		super().__init__()
		self.alpha = None

	def forward(self, x):
		if self.alpha is None:
			self.alpha = torch.nn.Parameter(torch.ones(x.shape[0]), requires_grad=True)
		...

采用这种写法的话，必须要在正式训练模型之前进行一次预推理，该预推理可以是伪输入数据的推理，目的是预推理时构建好每一层所需要的self.alpha可学习参数。我常用的写法如下：

dummy_input = torch.randn(1, 3, 32, 32)
# 1：batch size为1，只推理单个样本；3：数据集的图像通道数；32：数据集的图像大小
model(dummy_input)

必须要注意的是，新定义的self.alpha必须要放入optimizer中才可以训练，因此，上面这段预推理的代码必须要放在声明optimizer之前！！！原因很简单，声明optimizer时，有个传入参数就是模型参数列表：

optimizer = torch.optim.SGD(model.parameters(), xxx)

但是这里会出现一个问题：由于self.alpha时在模型运行（预推理）时构建的，所以尚未放入cuda中。因此，需要手动将self.alpha放入cuda中。于是，有如下两种可能的写法：

# 写法1（错误）
self.alpha = torch.nn.Parameter(torch.ones(x.shape[0]), requires_grad=True).to(x.device)

# 写法2（正确）
self.alpha = torch.nn.Parameter(torch.ones(x.shape[0]).to(x.device), requires_grad=True)