To be more specific, I have followed the ImageNet example in Nvidia Apex. I write the loss function inside my Network module which is like the follows:
class Network(torch.nn.Module):
    """
    Wrap the network module as well as the loss module on all GPUs to balance the computation among GPUs.
    """
    def __init__(self, opt, config, bn=False, dist=False):
        super(Network, self).__init__()
        self.posenet = PoseNet(opt.nstack, opt.hourglass_inp_dim, config.num_layers, bn=bn)
        # If we use train_parallel, we implement the parallel loss . And if we use train_distributed,
        # we should use single process loss because each process on these 4 GPUs  is independent
        self.criterion = MultiTaskLoss(opt, config) if dist else MultiTaskLossParallel(opt, config)

    def forward(self, inp_imgs, target_tuple):
        # Batch will be divided and Parallel Model will call this forward on every GPU
        output_tuple = self.posenet(inp_imgs)
        loss = self.criterion(output_tuple, target_tuple)

        if not self.training:
            # output will be concatenated  along batch channel automatically after the parallel model return
            return output_tuple, loss
        else:
            # output will be concatenated  along batch channel automatically after the parallel model return
            return loss

The training loss seems normal:




Epoch 0
train_loss: 589.6713480631511
val_loss: 536.4533081054688




Epoch 1
train_loss: 446.2322041829427
val_loss: 440.89935302734375


Epoch 2
train_loss: 436.07487325032554
val_loss: 433.20953369140625


Epoch 3
train_loss: 433.3325126139323
val_loss: 396.94744873046875


Epoch 4
train_loss: 425.1072373453776
val_loss: 406.3310546875


Epoch 5
train_loss: 418.57773783365883
val_loss: 392.5045166015625


Epoch 6
train_loss: 409.60796936035155
val_loss: 419.2001037597656


Epoch 7
train_loss: 410.79097737630207
val_loss: 409.8291320800781


Epoch 8
train_loss: 404.4842706298828
val_loss: 407.05352783203125


Epoch 9
train_loss: 399.4785394287109
val_loss: 388.7215881347656


Epoch 10
train_loss: 389.387607421875
val_loss: 379.6018981933594


Epoch 11
train_loss: 386.5943516031901
val_loss: 397.2137451171875


Epoch 12
train_loss: 382.25890686035154
val_loss: 376.7177734375


Epoch 13
train_loss: 387.2037613932292
val_loss: 360.4934387207031


Epoch 14
train_loss: 379.99100199381513
val_loss: 377.1543884277344


Epoch 15
train_loss: 381.0046073404948
val_loss: 378.36041259765625


Epoch 16
train_loss: 378.6185076904297
val_loss: 365.29205322265625


Epoch 17
train_loss: 380.5766967773437
val_loss: 364.39569091796875


Epoch 18
train_loss: 382.2865834554037
val_loss: 368.50152587890625



But the model seems not to have been trained well and the prediction results refuse to get better （which is bad actually).
I have struggled with this problem for a while. If I don’t use distributed training or Apex auto mixed-precision and I only wrap my Network module with torch.nn.parallel.DataParallel, everything goes fine and the prediction is good.