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)

试问这两种写法都正确吗？思考一分钟…
时间到！实际上，只有写法2是正确的！

写法1先定义nn.Parameter，后放入cuda，会导致参数重新变回到tensor，从而不可学习；
写法2先放入cuda，后定义nn.Parameter，可以成功定义参数，可以学习。

总之，记住就好，这确实也是一个一找可以找一整晚的BUG了