Repository: ludwig-ai/ludwig

Issue:
gpu_memory_limit only works if gpus parameter is specified
I'm trying to use the new gpu_memory_limit parameter introduced by TF2 as opposed as the gpu_fraction one.
One thing I noticed is that it only works if `--gpus` is also specified, as per
https://github.com/uber/ludwig/blob/b31ce876f6cd364471f7154740485656ab4ea6c9/ludwig/utils/tf_utils.py#L87-L107
Once you both specify `--gpus 0 --gpu_memory_limit 4000` (for example), allocated GPU memory is correctly set to 4GB. Without the `--gpus 0` bit, all available memory is allocated as if no `--gpu_memory_limit 4000` was specified.

One more thing is that Hyperopt still references the "old" `--gpu_fraction`
https://github.com/uber/ludwig/blob/b31ce876f6cd364471f7154740485656ab4ea6c9/ludwig/hyperopt_cli.py#L473-L479


Discussion:
Hey @carlogrisetti, thanks for reporting this issue.  I'll take a look at this today.

Resolving patch:
diff --git a/ludwig/hyperopt_cli.py b/ludwig/hyperopt_cli.py
--- a/ludwig/hyperopt_cli.py
+++ b/ludwig/hyperopt_cli.py
@@ -471,11 +471,11 @@ def cli(sys_argv):
         help="list of gpus to use"
     )
     parser.add_argument(
-        "-gf",
-        "--gpu_fraction",
-        type=float,
-        default=1.0,
-        help="fraction of gpu memory to initialize the process with",
+        '-gml',
+        '--gpu_memory_limit',
+        type=int,
+        default=None,
+        help='maximum memory in MB to allocate per GPU device'
     )
     parser.add_argument(
         "-uh",
diff --git a/ludwig/utils/tf_utils.py b/ludwig/utils/tf_utils.py
--- a/ludwig/utils/tf_utils.py
+++ b/ludwig/utils/tf_utils.py
@@ -84,27 +84,26 @@ def initialize_tensorflow(gpus=None,
         gpus = gpus.strip()
         gpus = [int(g) for g in gpus.split(",")]
 
-    if gpus:
-        if len(gpus) == 1 and gpus[0] == -1:
-            # CUDA_VISIBLE_DEVICES syntax for disabling all GPUs
-            tf.config.set_visible_devices([], 'GPU')
-        else:
-            # Allow memory growth and set memory limit. Regardless of whether we do this
-            # before or after setting visible devices, TensorFlow will allocate a small
-            # amount of memory per device.
-            for gpu in gpu_devices:
-                tf.config.experimental.set_memory_growth(gpu, True)
-                if gpu_memory_limit is not None:
-                    tf.config.set_logical_device_configuration(
-                        gpu,
-                        [tf.config.LogicalDeviceConfiguration(
-                            memory_limit=gpu_memory_limit)])
-
-            # Set visible devices so GPU utilization is isolated
-            # (no GPU contention between workers).
-            if gpu_devices:
-                local_devices = [gpu_devices[g] for g in gpus]
-                tf.config.set_visible_devices(local_devices, 'GPU')
+    if gpus and len(gpus) == 1 and gpus[0] == -1:
+        # CUDA_VISIBLE_DEVICES syntax for disabling all GPUs
+        tf.config.set_visible_devices([], 'GPU')
+    else:
+        # Allow memory growth and set memory limit. Regardless of whether we do this
+        # before or after setting visible devices, TensorFlow will allocate a small
+        # amount of memory per device.
+        for gpu in gpu_devices:
+            tf.config.experimental.set_memory_growth(gpu, True)
+            if gpu_memory_limit is not None:
+                tf.config.set_logical_device_configuration(
+                    gpu,
+                    [tf.config.LogicalDeviceConfiguration(
+                        memory_limit=gpu_memory_limit)])
+
+        # Set visible devices so GPU utilization is isolated
+        # (no GPU contention between workers).
+        if gpus and gpu_devices:
+            local_devices = [gpu_devices[g] for g in gpus]
+            tf.config.set_visible_devices(local_devices, 'GPU')
 
     _set_tf_init_params(param_tuple)
 


Test patch: