eml.session

distribute_config

To properly configure GPU memory usage and enable high-speed GPU-to-GPU communication always use eml.session.distribute_config when creating a TensorFlow session, if you already have a config.

import engineml.tensorflow as eml
import tensorflow as tf
config = tf.ConfigProto()
config = eml.session.distribute_config(config)
with tf.Session(config=config) as sess:
...

init_devices

To pin your PyTorch code to a single GPU and to enable high-speed GPU-to-GPU communication, use eml.session.init_devices. Always be sure to run eml.session.init_devices before any .cuda() calls.

import engineml.torch as eml
eml.session.init_devices()
model = MyModel().cuda()

init_op

After initializing global variables, use eml.session.init_op() to synchronize the model replica weights.

import engineml.tensorflow as eml
import tensorflow as tf
with tf.Session(config=config) as sess:
sess.run(tf.global_variables_initializer())
sess.run(eml.session.init_op())
...

init_op_hook

If you are using hooks to initialize your variables, use eml.session.init_op_hook() to synchronize the model replica weights before you begin training.

import engineml.tensorflow as eml
import tensorflow as tf
hooks = eml.session.init_op_hook()
with tf.train.MonitoredTrainingSession(..., hooks=hooks) as mon_sess:
...

make_distributed_config

To properly configure GPU memory usage and enable high-speed GPU-to-GPU communication always use eml.session.make_distributed_config() when creating a TensorFlow session, if you haven't already created a config.

import engineml.tensorflow as eml
import tensorflow as tf
with tf.Session(config=eml.session.make_distributed_config()) as sess:
...

sync_model_replicas

To synchronize your PyTorch model replica weights before training, use eml.session.sync_model_replicas.

import engineml.torch as eml
import torch
eml.session.init_devices()
model = MyModel().cuda()
lr = 0.1
lr = eml.optimizer.scale_learning_rate(lr)
opt = torch.optim.Adam(model.parameters(), lr=lr)
opt = eml.optimizer.distribute(opt)
eml.session.sync_model_replicas(model)