nn.DataParallel doesn’t work with a batch size of 1, but what about the functional version of DataParallel, data_parallel? data_parallel seems like it might be the PyTorch version of Lua/Torch7’s nn.GPU()
The functions themselves even seem to have basically the same inputs:
def data_parallel(module, inputs, device_ids=None, output_device=None, dim=0, module_kwargs=None):

function GPU:__init(module, device, outdevice)

Would something like this (put inside a loop) work better than just putting if statements everywhere for multiple gpus? Each nn.data_parallel module would have 1 “input gpu” and 1 “output gpu”.
net.add_module(layer_name, nn.data_parallel(layers[i], gpus[i], out_device))

Source:

  
      github.com
  
  
    torch/nn/blob/master/GPU.lua
------------------------------------------------------------------------
--[[ GPU ]]--
-- Decorates a module such that its parameters are
-- hosted on a specified GPU device.
-- The operations are also executed on that device.
-- Arguments input and gradOutput are converted to the specified device 
-- before being fed to the decorated module. 
-- Returned output is on the specified outdevice (defaults to device). 
-- Returned gradInput is allocated on the same device as the input.
-- The unit test is located in cunn.
------------------------------------------------------------------------
local GPU, parent = torch.class("nn.GPU", "nn.Container")

function GPU:__init(module, device, outdevice)
   parent.__init(self)
   assert(torch.type(device) == 'number')
   self.device = device
   self.outdevice = outdevice or device
   
   assert(torch.isTypeOf(module, 'nn.Module'))


  This file has been truncated. show original

  
  
    
    
  
  



  
      github.com
  
  
    pytorch/pytorch/blob/master/torch/nn/parallel/data_parallel.py#L133-L163 1

def data_parallel(module, inputs, device_ids=None, output_device=None, dim=0, module_kwargs=None):
r"""Evaluates module(input) in parallel across the GPUs given in device_ids.


This is the functional version of the DataParallel module.


Args:
    module: the module to evaluate in parallel
    inputs: inputs to the module
    device_ids: GPU ids on which to replicate module
    output_device: GPU location of the output  Use -1 to indicate the CPU.
        (default: device_ids[0])
Returns:
    a Tensor containing the result of module(input) located on
    output_device
"""
if not isinstance(inputs, tuple):
    inputs = (inputs,)


if device_ids is None:
    device_ids = list(range(torch.cuda.device_count()))


  This file has been truncated. show original