DataParallel модель обучения застрял на обратном проходе

Question

Я имею дело с проблемой PyTorch 1.1, где я тренирую модель на данных временных рядов.Данные хранятся в списке Tensors, где каждый тензор может быть разбит на несколько пакетов для распараллеливания.

Когда размер пакета равен 1, используется только один из двух графических процессоров (что ожидается), имодель тренируется плавно.Однако, когда размер пакета> 1 (это не переменная размера mini_batch, а количество пакетов, которое модель может использовать для распараллеливания), программа застревает при обратном распространении (line loss.backward ()).При отладке переменная потерь идентична (за исключением самого значения с плавающей запятой) по своим свойствам как в многопартийном, так и в однопакетном случаях, поэтому я не могу понять, что происходит неправильно.

Я имеюобнаружил, что у PyTorch были проблемы с графическими процессорами Tesla K80, но я также попробовал с Tesla P100s, и проблема сохраняется, поэтому проблема, похоже, не возникает.Я также просмотрел почти все связанные посты здесь, но ни одна из них, похоже, не решает / решает проблему, которая у меня сейчас есть.Мой учебный код прилагается ниже.

Любая помощь очень ценится.Спасибо!

def train_minibatch(self, model, optimizer, mini_batch=10):
        model.train()

        train_loss_value = 0
        predictions = []
        targets = []

        for batch_idx, ts in enumerate(self.train_loader):

            batch_loss_value = 0
            batch_predictions = []
            batch_targets = []

            optimizer.zero_grad()

            for i in range(len(ts)):

                output = model(ts[i].to(self.device))
                target = torch.argmax(cues[:, i], dim=1)
                prediction = output.max(1, keepdim=True)

                torch.cuda.synchronize()

                loss = sum([F.nll_loss(output[j, :].unsqueeze(0), target[j].unsqueeze(-1), weight=self.w) for j in range(len(target))])
                loss = loss / mini_batch
                loss.backward()

                if i > 0 and i % mini_batch == 0:
                    optimizer.step()
                    optimizer.zero_grad()

                train_loss_value += loss.item()
                predictions.extend(prediction.tolist())
                targets.extend(target.tolist())

                batch_loss_value += loss.item()
                batch_predictions.extend(prediction.tolist())
                batch_targets.extend(target.tolist())

            self.print_eval(batch_loss_value, batch_predictions, batch_targets, idx=batch_idx, header='batch idx:')