Я хочу модель PyTorch, которая условно использует разные слои

#pytorch

#pytorch

Вопрос:

Извините, если этот вопрос где-то там, я, кажется, ничего не могу понять.

Я знаю, что этого можно достичь, используя отдельные модели, но я хотел бы иметь единую модель, в которой используются разные слои в зависимости от входных данных. Что-то вроде:

 def __init__(self, n_inputs, n_outputs_arr):
    self.linear = []
    for n_outputs in n_outputs_arr:
        self.linear.append(nn.Linear(n_inputs, n_outputs))

def forward(self, x, i):
    return self.linear[i](x)
 

и, конечно, обновлять веса только на i ‘-м уровне (пакеты будут состоять только из одного i ). Это что?

Комментарии:

1. Показанный код должен работать, в чем проблема?

2. не используйте python list для nn объектов. используйте nn.ModuleList . проверьте это

3. Спасибо, Гириш, это сработало отлично!