smth:

in forward:
scatter mini-batch to GPU1, GPU2
replicate model on GPU2 (it is already on GPU1)
model_gpu1(input_gpu1), model_gpu2(input_gpu2) (this step is parallel_apply)
gather output mini-batch from GPU1, GPU2 onto GPU1
in backward:- scatter grad_output and input- parallel_apply model’s backward pass- reduce GPU2 replica’s gradients onto GPU1 model- ## Now there is only a single model again with accumulated gradients from GPU1 and GPU2- gather the grad_input


Soumith and Adam, I am having a great time exploring PyTorch!  Thanks for the awesome library.
I am trying to saturate a 64-core/256-thread CPU in addition to the GPUs.  Any pointers on how I can extend Data_parallel.py to create 3 scatters on GPU0, GPU1, and CPU(0-255)?
With Keras I modified this script to saturate the CPU:

  
      github.com
  
  
    kuza55/keras-extras/blob/master/utils/multi_gpu.py 14
from keras.layers import merge
from keras.layers.core import Lambda
from keras.models import Model

import tensorflow as tf

def make_parallel(model, gpu_count):
    def get_slice(data, idx, parts):
        shape = tf.shape(data)
        size = tf.concat([ shape[:1] // parts, shape[1:] ],axis=0)
        stride = tf.concat([ shape[:1] // parts, shape[1:]*0 ],axis=0)
        start = stride * idx
        return tf.slice(data, start, size)

    outputs_all = []
    for i in range(len(model.outputs)):
        outputs_all.append([])

    #Place a copy of the model on each GPU, each getting a slice of the batch
    for i in range(gpu_count):


  This file has been truncated. show original