Working with different backends¶
DGL supports PyTorch, MXNet and Tensorflow backends. To change them, set the
environcment variable. The default backend is PyTorch.
pytorch to specify PyTorch backend. The required PyTorch
version is 0.4.1 or later. See pytorch.org for installation instructions.
mxnet to specify MXNet backend. The required MXNet version is
1.5 or later. See mxnet.apache.org for installation
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
flag. See this FAQ for more information.
MXNet 1.5 and later has an option to enable Numpy shape mode for
NDArray objects, some DGL models
need this mode to be enabled to run correctly. However, this mode may not compatible with pretrained
model parameters with this mode disabled, e.g. pretrained models from GluonCV and GluonNLP.
DGL_MXNET_SET_NP_SHAPE, users can switch this mode on or off.
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
true to prevent Tensorflow take over the whole GPU memory:
pip install tfdlpack # when using tensorflow cpu version
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