Запускать несколько моделей ансамбля параллельно с PyTorch

Question

Моя нейронная сеть имеет следующую архитектуру:

input -> 128x (separate fully connected layers) -> output averaging

Я использую ModuleList для хранения списка полностью подключенных слоев. Вот как это выглядит в этой точке:

class MultiHead(nn.Module):
    def __init__(self, dim_state, dim_action, hidden_size=32, nb_heads=1):
        super(MultiHead, self).__init__()

        self.networks = nn.ModuleList()
        for _ in range(nb_heads):
            network = nn.Sequential(
                nn.Linear(dim_state, hidden_size),
                nn.Tanh(),
                nn.Linear(hidden_size, dim_action)
            )
            self.networks.append(network)

        self.cuda()
        self.optimizer = optim.Adam(self.parameters())

Затем, когда мне нужно вычислить вывод, я использую конструкцию for ... in для выполнения прямого и обратного прохождения через все слои:

q_values = torch.cat([net(observations) for net in self.networks])

# skipped code which ultimately computes the loss I need

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

Это работает! Но мне интересно, не смогу ли я сделать это более эффективно. Я чувствую, что, выполняя for...in, я на самом деле прохожу каждый отдельный слой FC один за другим, хотя я ожидаю, что эта операция может выполняться параллельно.

Answer 1

В случае Convnd вместо Linear вы могли бы использовать аргумент groups для "сгруппированных сверток" (или "глубинных сверток"). Это позволяет вам обрабатывать все параллельные сети одновременно.

Если вы используете сверточное ядро ​​размером 1, то свертка ничего не делает, кроме применения слоя Linear, где каждый канал считается входным измерением. Таким образом, грубая структура вашей сети будет выглядеть следующим образом:

1. Измените входной тензор формы B x dim_state следующим образом: добавьте дополнительное измерение и выполните репликацию nb_state -times B x dim_state для B x (dim_state * nb_heads) x 1
2. заменить два Linear на

nn.Conv1d(in_channels=dim_state * nb_heads, out_channels=hidden_size * nb_heads, kernel_size=1, groups=nb_heads)

и

nn.Conv1d(in_channels=hidden_size * nb_heads, out_channels=dim_action * nb_heads, kernel_size=1, groups=nb_heads)

теперь у нас есть тензор размером B x (dim_action x nb_heads) x 1, теперь вы можете изменить его на любую нужную вам форму (например, B x nb_heads x dim_action)

---

Хотя CUDA изначально поддерживает сгруппированные свертки, былонекоторые проблемы в pytorch со скоростью сгруппированных извилин (см., например, здесь ), но я думаю, что это было решено сейчас.