Working with different backends

DGL supports PyTorch, MXNet and Tensorflow backends. To change them, set the DGLBACKEND environcment variable. The default backend is PyTorch.

PyTorch backend

Export DGLBACKEND as pytorch to specify PyTorch backend. The required PyTorch version is 0.4.1 or later. See pytorch.org for installation instructions.

MXNet backend

Export DGLBACKEND as mxnet to specify MXNet backend. The required MXNet version is 1.5 or later. See mxnet.apache.org for installation instructions.

MXNet uses uint32 as the default data type for integer tensors, which only supports graph of size smaller than 2^32. To enable large graph training, build MXNet with USE_INT64_TENSOR_SIZE=1 flag. See this FAQ for more information.

Tensorflow backend

Export DGLBACKEND as tensorflow to specify Tensorflow backend. The required Tensorflow version is 2.0 or later. See tensorflow.org for installation instructions. In addition, Tensorflow backend requires tfdlpack package installed as follows and set TF_FORCE_GPU_ALLOW_GROWTH to true to prevent Tensorflow take over the whole GPU memory:

pip install tfdlpack  # when using tensorflow cpu version

or

pip install tfdlpack-gpu  # when using tensorflow gpu version
export TF_FORCE_GPU_ALLOW_GROWTH=true # and add this to your .bashrc/.zshrc file if needed