python - PyTorch : How to properly create a list of nn. 线性()

Question

我创建了一个以 nn.Module 作为子类的类。

在我的类(class)中，我必须创建 N 个线性变换，其中 N 作为类参数给出。

因此我进行如下操作:

    self.list_1 = []

    for i in range(N):
        self.list_1.append(nn.Linear(self.x, 1, bias=mlp_bias))

在 forward 方法中，我调用这些矩阵(使用 list_1[i])并连接结果。

两件事:

1)

即使我使用 model.cuda()，这些线性变换在 cpu 上使用，我得到以下错误:

RuntimeError: Expected object of type Variable[torch.cuda.FloatTensor] but found type Variable[torch.FloatTensor] for argument #1 'mat2'

我必须做

self.list_1.append(nn.Linear(self.x, 1, bias=mlp_bias).cuda())

这不是必需的，如果我这样做的话:

self.nn = nn.Linear(self.x, 1, bias=mlp_bias)

然后直接使用self.nn。

2)

由于更明显的原因，当我在 main 中打印(模型)时，列表中的线性矩阵没有打印出来。

还有没有别的办法。也许使用 bmm ？我发现它不太容易，实际上我想分别获得我的 N 个结果。

提前谢谢你，

中号

Answer 1

您可以使用 nn.ModuleList 来包装您的线性层列表，如 here 所述

self.list_1 = nn.ModuleList(self.list_1)