muji Namespace Reference

Module caffe2.python.muji. More...

Functions
def	OnGPU (gpu_id)
def	OnCPU ()
def	Allreduce (net, blobs, reduced_affix="_reduced", gpu_indices=None)
def	Allreduce2 (net, blobs, reduced_affix, gpu_indices)
def	Allreduce4 (net, blobs, reduced_affix, gpu_indices)
def	Allreduce8 (net, blobs, reduced_affix, gpu_indices)
def	AllreduceFallback (net, blobs, reduced_affix, gpu_indices)

Detailed Description

Module caffe2.python.muji.

Function Documentation

Allreduce()

def muji.Allreduce	(		net,
			blobs,
			reduced_affix = "_reduced",
			gpu_indices = None
	)

The general Allreduce interface that reroutes the function calls.

Definition at line 33 of file muji.py.

Allreduce2()

def muji.Allreduce2	(		net,
			blobs,
			reduced_affix,
			gpu_indices
	)

Allreduce for 2 gpus.

  Algorithm: 0r <- 0 + 1, 1r <- 0r, where r means "reduced"

Definition at line 53 of file muji.py.

Allreduce4()

def muji.Allreduce4	(		net,
			blobs,
			reduced_affix,
			gpu_indices
	)

Allreduce for 4 gpus.

  Algorithm: 2 level reduction.
  0r <- 0 + 1, 2r <- 2 + 3
  0r <- 0r + 2r
  2r <- 0r,
  1r <- 0r, 3r <- 2r

Definition at line 69 of file muji.py.

Allreduce8()

def muji.Allreduce8	(		net,
			blobs,
			reduced_affix,
			gpu_indices
	)

Allreduce for 8 gpus.

  Algorithm: 3 level reduction.
  0r <- 0 + 1, 2r <- 2 + 3, 4r <- 4 + 5, 6r <- 6 + 7
  0r <- 0r + 2r, 4r <- 4r + 6r
  0r <- 0r + 4r
  4r <- 0r
  2r <- 0r, 6r <- 4r
  1r <- 0r, 3r <- 2r, 5r <- 4r, 7r <- 6r

Definition at line 109 of file muji.py.

AllreduceFallback()

def muji.AllreduceFallback	(		net,
			blobs,
			reduced_affix,
			gpu_indices
	)

A fallback option for Allreduce with no assumption on p2p.

  Algorithm: a flat operation on gpu 0
  0r <- 0
  0r <- 0r + i for i in gpu_indices[1:]
  ir <- 0r for i in gpu_indices[1:]

Definition at line 169 of file muji.py.

OnGPU()

def muji.OnGPU

(

gpu_id

)

A utility function that returns a device option protobuf of the
  specified gpu id.

Definition at line 17 of file muji.py.