Pytorch динамическое количество слоев?

Question

Я пытаюсь указать динамическое количество слоев, что, кажется, делаю неправильно. Моя проблема в том, что когда я определяю здесь 100 слоев, я получаю ошибку на шаге вперед. Но когда я правильно определяю слой, он работает? Ниже упрощенный пример

class PredictFromEmbeddParaSmall(LightningModule):
def __init__(self, hyperparams={'lr': 0.0001}):
    super(PredictFromEmbeddParaSmall, self).__init__()
    #Input is something like tensor.size=[768*100]
    self.TO_ILLUSTRATE = nn.Linear(768, 5)
    self.enc_ref=[]
    for i in range(100):
        self.enc_red.append(nn.Linear(768, 5))
    # gather the layers output sth
    self.dense_simple1 = nn.Linear(5*100, 2)
    self.output = nn.Sigmoid()
def forward(self, x):
    # first input to enc_red
    x_vecs = []
    for i in range(self.para_count):
        layer = self.enc_red[i]
        # The first dim is the batch size here, output is correct
        processed_slice = x[:, i * 768:(i + 1) * 768]
        # This works and give the out of size 5
        rand = self.TO_ILLUSTRATE(processed_slice)
        #This will fail? Error below
        ret = layer(processed_slice)
        #more things happening we can ignore right now since we fail earlier

Я получаю эту ошибку при выполнении "ret = layer.forward(processing_slice)"

RuntimeError: ожидаемый объект типа устройства cuda, но получил процессор типа устройства для аргумента #1 'self' при вызове _th_addmm

Есть ли более умный способ запрограммировать это? ИЛИ решить ошибку?

Answer 1

Вы должны использовать ModuleList из pytorch вместо списка: https://pytorch.org/docs/master/generated/torch.nn.ModuleList.html. Это связано с тем, что Pytorch должен поддерживать график со всеми модулями вашей модели, если вы просто добавите их в список, они не будут правильно проиндексированы на графике, что приведет к ошибке, с которой вы столкнулись.

Ваш coude должен быть чем-то похожим:

class PredictFromEmbeddParaSmall(LightningModule):
def __init__(self, hyperparams={'lr': 0.0001}):
    super(PredictFromEmbeddParaSmall, self).__init__()
    #Input is something like tensor.size=[768*100]
    self.TO_ILLUSTRATE = nn.Linear(768, 5)
    self.enc_ref=nn.ModuleList()                     # << MODIFIED LINE <<
    for i in range(100):
        self.enc_red.append(nn.Linear(768, 5))
    # gather the layers output sth
    self.dense_simple1 = nn.Linear(5*100, 2)
    self.output = nn.Sigmoid()
def forward(self, x):
    # first input to enc_red
    x_vecs = []
    for i in range(self.para_count):
        layer = self.enc_red[i]
        # The first dim is the batch size here, output is correct
        processed_slice = x[:, i * 768:(i + 1) * 768]
        # This works and give the out of size 5
        rand = self.TO_ILLUSTRATE(processed_slice)
        #This will fail? Error below
        ret = layer(processed_slice)
        #more things happening we can ignore right now since we fail earlier

Тогда все должно работать нормально!

Изменить: альтернативный способ.

Вместо того, чтобы использовать ModuleList вы также можете просто использовать nn.Sequential, это позволяет избежать использования forпетля в прямом проходе. Это также означает, что у вас не будет доступа к промежуточным активациям, так что это не решение для вас, если они вам нужны.

class PredictFromEmbeddParaSmall(LightningModule):
def __init__(self, hyperparams={'lr': 0.0001}):
    super(PredictFromEmbeddParaSmall, self).__init__()
    #Input is something like tensor.size=[768*100]
    self.TO_ILLUSTRATE = nn.Linear(768, 5)
    self.enc_ref=[]
    for i in range(100):
        self.enc_red.append(nn.Linear(768, 5))

    self.enc_red = nn.Seqential(*self.enc_ref)       # << MODIFIED LINE <<
    # gather the layers output sth
    self.dense_simple1 = nn.Linear(5*100, 2)
    self.output = nn.Sigmoid()
def forward(self, x):
    # first input to enc_red
    x_vecs = []
    out = self.enc_red(x)                            # << MODIFIED LINE <<

Answer 2

Немного более гибкое решение, которое зависит от вкуса или сложности конкретной ситуации, было опубликовано здесь .

Для справки я размещаю скорректированную версию кода здесь:

      import torch
from torch import nn, optim
from torch.nn.modules import Module
from implem.settings import settings


class Model(nn.Module):
    def __init__(self, input_size, layers_data: list, learning_rate=0.01, optimizer=optim.Adam):
        super().__init__()
        self.layers = nn.ModuleList()
        self.input_size = input_size  # Can be useful later ...
        for size, activation in layers_data:
            self.layers.append(nn.Linear(input_size, size))
            input_size = size  # For the next layer
            if activation is not None:
                assert isinstance(activation, Module), \
                    "Each tuples should contain a size (int) and a torch.nn.modules.Module."
                self.layers.append(activation)

        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.learning_rate = learning_rate
        self.optimizer = optimizer(params=self.parameters(), lr=learning_rate)

    def forward(self, input_data):
        for layer in self.layers:
            input_data = layer(input_data)
        return input_data


# test that the net is working properly 
if __name__ == "__main__":
    data_size = 5
    layer1, layer2 = 10, 10
    output_size = 2
    data = torch.randn(data_size)
    mlp = Model(data_size, [(layer1, nn.ReLU()), (layer2, nn.ReLU()), (output_size, nn.Sigmoid())])
    output = mlp(data)
    print("done")