merlin.dataloader.tf_utils.configure_tensorflow

merlin.dataloader.tf_utils.configure_tensorflow(memory_allocation=None, device=None)[source]

Control the GPU memory allocation that is performed when using TensorFlow.

Example usage:

# Allocate 20% of GPU memory to TensorFlow on the first
# GPU in the system.
configure_tensorflow(.2, 0)

Parameters

memory_allocation (float, optional) – Value between 0 and 1 that represents the fraction of GPU memory to allocate to TensorFlow. This parameter overrides the TF_MEMORY_ALLOCATION environment variable. If you do not specify a value and the environment variable is not set, 50% of GPU memory is allocated to TensorFlow. The default value is None.
device (Int, optional) – Integer representing the index of the GPU to run on. This parameter overrides the TF_VISIBLE_DEVICE environment variable. The default value is None.