PyTorch中定义可学习参数时的坑

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

class model(torch.nn.Module):def __init__(self):super().__init__()self.alpha = Nonedef 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)