Multi-GPU support

[ncullen93]

Multi-GPU support

[jph00]

I just noticed this in practice - tried to train using nn.DataParallel and got the error that the tensors were on the wrong GPU.

[achaiah]

Yes, it's not totally as trivial as just stuffing everything into nn.DataParallel but it's not terribly difficult either. I've done it in the past in my custom code but would love to switch to this framework as it has additional callbacks, transforms etc that are nice to use.

[achaiah]

@jph00 you may want to examine your code to make sure you're working with the same GPU. The parallelization using nn.DataParallel seems to work for me if I leave out regularizers and constraints. I'm still trying to understand why but it basically complains in module_trainer.py:

    593                         regularizer_loss = regularizers(self.model)
    594                         loss += regularizer_loss
--> 595                         batch_logs['regularizer_loss'] = regularizer_loss.data[0]

AttributeError: 'float' object has no attribute 'data'

Looking at the RegularizerModule I don't even see where data attribute is coming from. It seems that loss is always a float... yet somehow the code runs fine when executed on a single GPU. Maybe @ncullen93 has an idea. I'm sure the fix is simple.

P.S. there's a similar issue with constraints.

[achaiah]

I can get around the issue by simply commenting out the lines that store regularization and constraints loss in batch_logs but it's obviously not ideal. Not sure what the real solution is.

Also, there's an additional problem in the module_trainer.py summary:

torchsample/torchsample/modules/module_trainer.py in summary(self, input_size)
    113         self.model.apply(register_hook)
    114         # make a forward pass
--> 115         self.model(x)
    116         # remove these hooks
    117         for h in hooks:

TypeError: Broadcast function not implemented for CPU tensors

[achaiah]

I added a pull request partially addressing the multi-GPU issues: #40