torch.nn.DataParallel.forward()
# 将张量切成大致相等的块并将它们分布在给定的GPU上。 复制对不是张量的对象的引用。
inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids)
if len(self.device_ids) == 1:
return self.module(*inputs[0], **kwargs[0])
# 复制多份模型(包括子节点)到每个gpu,并返回为replicas
replicas = self.replicate(self.module, self.device_ids[:len(inputs)])
# 通过多线程调用output = module(*input, **kwargs)
outputs = self.parallel_apply(replicas, inputs, kwargs)
# 从不同gpu收集tensors
return self.gather(outputs, self.output_device)