Создание NN-модели Feed Forward в динамическом количестве скрытых слоев

Question

Почему эти два сегмента кода не эквивалентны: Сегмент 1: создание модели с 2 слоями.

class FNNModule(nn.Module):
    def __init__(self, input_dim, output_dim, hidden_dim1, hidden_dim2, non_linear_function):
        super().__init__()
        self.hidden1 = nn.Linear(input_dim, hidden_dim1)
        self.hidden2 = nn.Linear(hidden_dim1, hidden_dim2)       
        self.non_linear_function = non_linear_function()
        self.final_linear = nn.Linear(hidden_dim2, output_dim)

    def forward(self, x):
        out = self.hidden1(x)
        out = self.non_linear_function(out)
        out = self.hidden2(x)
        out = self.non_linear_function(out)        
        out = self.final_linear(out)
        return out

Сегмент два: создание той же модели, но с изменением кода, где hidden_layers является переменной:

class FNNModuleVar(nn.Module):
    def __init__(self, input_dim, output_dim, hidden_dim_array = [], non_linear_function_array=[]):
        super().__init__()
        self.linear_functions = []
        self.non_linear_functions = [x() for x in non_linear_function_array]
        self.hidden_layers = len(hidden_dim_array)
        for l in range(self.hidden_layers):
            self.linear_functions.append(nn.Linear(input_dim, hidden_dim_array[l]))
            input_dim = hidden_dim_array[l]
        self.final_linear = nn.Linear(input_dim, output_dim)

    def forward(self, x):
        out = x
        for i in range(self.hidden_layers):
            out = self.linear_functions[i](out)
            out = self.non_linear_functions[i](out)
        out = self.final_linear(x)
        return out
modelVar = FNNModuleVar(input_dim, output_dim, [100, 50], [nn.Tanh, nn.Tanh])
model = FNNModule(input_dim, output_dim, 100, 50, nn.Tanh)

Когда я пытаюсь выполнить итерации по modelVar.parameters() и model.parameters(), я вижу, что у меня очень разные модели.

Что я делаю неправильно в modelVar?

Answer 1

Эти модули вызываются так, как вы ожидаете, что они будут вызваны, они просто не видны для модуля. Чтобы сделать их видимыми, вы можете обернуть их в nn.ModuleList следующим образом:

class FNNModuleVar(nn.Module):
    def __init__(self, input_dim, output_dim, hidden_dim_array = [], non_linear_function_array=[]):
        super().__init__()
        self.linear_functions = []
        self.non_linear_functions = [x() for x in non_linear_function_array]
        self.hidden_layers = len(hidden_dim_array)
        for l in range(self.hidden_layers):
            self.linear_functions.append(nn.Linear(input_dim, hidden_dim_array[l]))
            input_dim = hidden_dim_array[l]
        self.linear_functions = nn.ModuleList(self.linear_functions)
        self.final_linear = nn.Linear(input_dim, output_dim)

    def forward(self, x):
        out = x
        for i in range(self.hidden_layers):
            out = self.linear_functions[i](out)
            out = self.non_linear_functions[i](out)
        out = self.final_linear(out)
        return out

печать моделей теперь даст:

FNNModule(
  (hidden1): Linear(in_features=50, out_features=100, bias=True)
  (hidden2): Linear(in_features=100, out_features=50, bias=True)
  (non_linear_function): Tanh()
  (final_linear): Linear(in_features=50, out_features=100, bias=True)
)
FNNModuleVar(
  (linear_functions): ModuleList(
    (0): Linear(in_features=50, out_features=100, bias=True)
    (1): Linear(in_features=100, out_features=50, bias=True)
  )
  (final_linear): Linear(in_features=50, out_features=100, bias=True)
)

Подробнее: https://pytorch.org/docs/stable/nn.html#torch.nn.ModuleList