在Python列表上使用nn.ModuleList會大大減慢訓練速度

Question

我正在訓練一個非常簡單的模型，該模型將隱藏層的數量作為參數。 我最初將這些隱藏層存儲在香草python列表[] ，但是當將此列表轉換為nn.ModuleList ，訓練速度至少會降低一個數量級 ！

AdderNet

class AdderNet(nn.Module):
    def __init__(self, num_hidden, hidden_width):
        super(AdderNet, self).__init__()
        self.relu = nn.ReLU()

        self.hiddenLayers = []
        self.inputLayer = nn.Linear(2, hidden_width)
        self.outputLayer = nn.Linear(hidden_width, 1)

        for i in range(num_hidden):
            self.hiddenLayers.append(nn.Linear(hidden_width, hidden_width))

        self.hiddenLayers = nn.ModuleList(self.hiddenLayers)  # <--- causes DRAMATIC slowdown!

    def forward(self, x):
        out = self.inputLayer(x)
        out = self.relu(out)

        for layer in self.hiddenLayers:
            out = layer(out)
            out = self.relu(out)

        return self.outputLayer(out)

訓練

for epoch in range(num_epochs):
    for i in range(0,len(data)):
        out = model.forward(data[i].x)
        loss = lossFunction(out, data[i].y)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

Answer 1

這是因為在使用普通的python列表時，參數不會添加到模型的參數列表中，而在使用ModuleList時會添加。 因此，在原始方案中，您從未真正訓練過隱藏層，這就是為什么它更快的原因。 （在每種情況下打印出model.parameters（），看看會發生什么！）