Add flag for multiple GPUs. User must manually support this.
@@ -125,12 +125,20 @@ class TorchAgent(Agent):
'text and __p2__ in front of past labels when available or '
'past utterances generated by the model. these are added to '
'the dictionary during initialization.')
- agent.add_argument(
+
+ # these gpu options are all mutually exclusive, and should error if the
+ # user tries to present multiple of them
+ gpugroup = agent.add_mutually_exclusive_group()
+ gpugroup.add_argument(
+ '-gpu', '--gpu', type=int, default=-1, help='which GPU to use')
+ gpugroup.add_argument(
+ '--multigpu', type='bool', default=False,
+ help='Suggest model should employ multiple GPUs. Not all models '
+ 'respect this flag.')
+ gpugroup.add_argument(
'--no-cuda', type='bool', default=False,
help='disable GPUs even if available. otherwise, will use GPUs if '
'available on the device.')
- agent.add_argument(
- '-gpu', '--gpu', type=int, default=-1, help='which GPU to use')
def __init__(self, opt, shared=None):
"""Initialize agent."""
@@ -155,9 +163,15 @@ class TorchAgent(Agent):
# check for cuda
self.use_cuda = not opt['no_cuda'] and torch.cuda.is_available()
if self.use_cuda:
+ if opt['multigpu'] and opt['gpu'] != -1:
+ raise ValueError("Can't use --multigpu and --gpu together.")
if not shared:
print('[ Using CUDA ]')
- torch.cuda.device(opt['gpu'])
+ if not shared and opt['gpu'] != -1:
+ torch.cuda.set_device(opt['gpu'])
+ # It's up to the user of TorchAgent to handle multiple GPU issues.
+ # A typical approach would be to employ nn.DataParallel
+ # (https://pytorch.org/docs/0.4.1/nn.html?highlight=dataparallel#torch.nn.DataParallel)
# now set up any fields that all instances may need
self.EMPTY = torch.Tensor([])
