Выполнение другой модели параллельно прямому проходу модели с помощью PyTorch

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

Я хочу иметь возможность выполнять вспомогательные вычисления после вычисления блока параллельно вычислению следующего блока ResNet, чтобы уменьшить сквозную задержку выполнения всего конвейера на ГПУ.

У меня есть базовый код, который работает корректно с точки зрения функциональности, но выполнение вспомогательной модели происходит последовательно по отношению к вычислению блока ResNet. Я проверил это двумя способами -

1. Добавляя операторы печати и проверяя порядок выполнения.

2. Инструментируя время работы исходной модели ResNet (скажем, время t1) и вспомогательной модели (скажем, время t2). Мое время выполнения в настоящее время t1+t2.

Исходный код блока ResNet (это BasicBlock, так как я экспериментирую с ResNet-18). Весь код доступен здесь

  class BasicBlock(nn.module):
  ...

      def forward(self, x):

        residual = x
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        out = self.conv2(out)
        out = self.bn2(out)

        if self.downsample is not None:
          residual = self.downsample(x)

        out += residual
        out = self.relu(out)
        return out

Это моя модификация, которая работает серийно -

def forward(self, x):
    if len(x[0]) == self.auxiliary_prediction_size: # Got an Auxiliary prediction earlier
        return x

    # Do usual block computation
    residual = x
    out = self.conv1(x)
    out = self.bn1(out)
    out = self.relu(out)
    out = self.conv2(out)
    out = self.bn2(out)

    if self.downsample is not None:
        residual = self.downsample(x)

    out += residual
    out = self.relu(out)

    # Try to make an auxiliary prediction
    # First flatten the tensor (also assume for now that batch size is 1)
    batchSize = x.shape[0]
    intermediate_output = out.view(batchSize, -1)
    # Place the flattened on GPU
    device = torch.device("cuda:0")
    input = intermediate_output.to(device)
    # Make auxiliary prediction
    auxiliary_input = out.float()
    auxiliary_prediction = self.auxiliary_model(auxiliary_input)
    if auxiliary_prediction meets some condition:
      return auxiliary_prediction

    # If no auxiliary prediction, then return intermediate output
    return out

Понятно, что приведенный выше код вызывает зависимость данных между выполнением вспомогательной модели и следующим блоком, и, следовательно, все происходит последовательно. Первое решение, которое я попробовал, состояло в том, чтобы проверить, уменьшает ли задержка разрыв этой зависимости данных. Я попытался сделать это, разрешив выполнение вспомогательной модели, но не возвращая вспомогательное_предсказание, если условие выполнено (обратите внимание, что это нарушило бы функциональность, но этот эксперимент был проведен исключительно для понимания поведения). По сути, то, что я сделал, было -

    batchSize = x.shape[0]
    intermediate_output = out.view(batchSize, -1)
    # Place the flattened on GPU
    device = torch.device("cuda:0")
    input = intermediate_output.to(device)
    # Make auxiliary prediction
    auxiliary_input = out.float()
    auxiliary_prediction = self.auxiliary_model(auxiliary_input)
    if auxiliary_prediction meets some condition:
      # Comment out return to break data dependency
      #return auxiliary_prediction

    # If no auxiliary prediction, then return intermediate output
    return out

Однако это не сработало, и после дальнейших исследований я наткнулся на потоки CUDA по адресу Ссылка на переполнение стека. Я попытался внедрить идею потоков CUDA, чтобы решить мою проблему следующим образом:

def forward(self, x):
    if len(x[0]) == self.auxiliary_prediction_size: # Got an Auxiliary prediction earlier
        return x

    s1 = torch.cuda.Stream()
    s2 = torch.cuda.Stream()

    with torch.cuda.Stream(s1):
      # Do usual block computation
      residual = x
      out = self.conv1(x)
      out = self.bn1(out)
      out = self.relu(out)
      out = self.conv2(out)
      out = self.bn2(out)

      if self.downsample is not None:
        residual = self.downsample(x)

      out += residual
      out = self.relu(out)

    with torch.cuda.Stream(s2):
      # Try to make an auxiliary prediction
      # First flatten the tensor (also assume for now that batch size is 1)
      out_detach = out.detach() # Detach from backprop flow and from computational graph dependency
      batchSize = x.shape[0]
      intermediate_output = out_detach.view(batchSize, -1)
      # Place the flattened on GPU
      device = torch.device("cuda:0")
      input = intermediate_output.to(device)
      # Make auxiliary prediction
      auxiliary_input = out_detach.float()
      auxiliary_prediction = self.auxiliary_model(auxiliary_input)
      if auxiliary_prediction meets some condition:
        return auxiliary_prediction

    # If no auxiliary prediction, then return intermediate output
    return out

Однако выходные данные Nvidia Visual Profiler по-прежнему указывают на то, что вся работа по-прежнему выполняется в потоке по умолчанию и по-прежнему сериализуется. Обратите внимание, что я проверил с помощью небольшой программы CUDA, что потоки CUDA поддерживаются версией CUDA, которую я использую.

Мои вопросы -

1. Почему нарушение зависимости от данных не приводит к тому, что PyTorch планирует параллельные вычисления? Я думал, что в этом суть графов динамических вычислений в PyTorch.

2. Почему при использовании потоков CUDA вычисления не делегируются потокам по умолчанию?

3. Существуют ли альтернативные подходы к выполнению вспомогательной модели асинхронно/параллельно расчету блока ResNet?