Introduction

In Python, you usually do not request “two GPUs” through one generic setting. What you actually control is which physical GPUs the process is allowed to see, and then how the ML framework uses those visible devices. The standard mechanism is CUDA_VISIBLE_DEVICES, followed by framework-specific configuration for single-device or multi-device execution.

Restrict Visible Devices First

The most common control point is the environment variable CUDA_VISIBLE_DEVICES. It tells the CUDA runtime which physical device IDs the process can see.

From the shell:

export CUDA_VISIBLE_DEVICES=0,1
python train.py

This means only physical GPU 0 and GPU 1 are visible to the process. If you set:

export CUDA_VISIBLE_DEVICES=3
python train.py

many frameworks will expose that one visible GPU to your code as local device 0. That renumbering is normal.

If You Set It in Python, Do It Early

Sometimes a small script sets the visibility inside Python. That can work, but only if it happens before importing a CUDA-backed framework.

1import os
2
3os.environ["CUDA_VISIBLE_DEVICES"] = "1"
4
5import torch
6
7print(torch.cuda.device_count())

If the framework has already initialized CUDA, changing the environment variable is usually too late. In notebooks, this often means you must restart the kernel.

TensorFlow Example

TensorFlow lets you inspect physical devices and restrict which ones remain visible to TensorFlow itself.

1import tensorflow as tf
2
3physical_gpus = tf.config.list_physical_devices("GPU")
4print("physical GPUs:", physical_gpus)
5
6if physical_gpus:
7    tf.config.set_visible_devices(physical_gpus[:1], "GPU")
8    tf.config.experimental.set_memory_growth(physical_gpus[0], True)
9
10visible_gpus = tf.config.get_visible_devices("GPU")
11print("visible GPUs:", visible_gpus)

This gives programmatic control, but it still needs to happen early in process startup.

If you want multi-GPU training in TensorFlow, making several devices visible is only the first step. You also need an explicit distribution strategy.

1import tensorflow as tf
2
3strategy = tf.distribute.MirroredStrategy()
4print("replicas:", strategy.num_replicas_in_sync)

Without that strategy, the model may still train on only one device.

PyTorch Example

PyTorch generally follows the environment-level visibility rules and then lets you choose local device indexes in code.

1import torch
2
3print("CUDA available:", torch.cuda.is_available())
4print("visible device count:", torch.cuda.device_count())
5
6if torch.cuda.is_available():
7    for i in range(torch.cuda.device_count()):
8        print(i, torch.cuda.get_device_name(i))
9
10    x = torch.tensor([1.0, 2.0, 3.0], device="cuda:0")
11    print(x)

This is enough to verify that the visibility settings worked. For real multi-GPU training, you still need an explicit training strategy such as DistributedDataParallel.

Visibility Is Not the Same as Parallelism

A very common misunderstanding is to make two GPUs visible and assume the framework will automatically split work across both. Visibility only defines which devices are available. It does not define how the computation is scheduled.

That is why the operational sequence should be:

1. restrict visibility
2. print device count and names
3. configure the training strategy explicitly

If step three never happens, the application may still run on a single GPU.

Process-Local Device Numbering

Suppose a machine has four GPUs and you launch with CUDA_VISIBLE_DEVICES=2,3. Inside the process, many libraries expose those two devices as local indexes 0 and 1.

This often confuses people debugging logs, but it is intentional. The process sees a filtered device list, not the host's global numbering directly.

A Good Operational Pattern

For shared servers and CI jobs, the cleanest pattern is usually to set visibility outside the script and let the Python code only verify and log what it sees. That keeps scheduler policy and application logic separate.

For example, on a shared machine you might let the scheduler set the environment and have the script print the discovered devices at startup. That makes debugging much easier than mixing resource allocation policy into the model code.

Common Pitfalls

• Treating CUDA_VISIBLE_DEVICES as a count instead of a list of device identifiers.
• Setting visibility after importing TensorFlow, PyTorch, or another CUDA-backed library.
• Assuming visible GPUs automatically imply multi-GPU execution.
• Forgetting that visible devices are often renumbered inside the process.
• Debugging GPU allocation without printing the actual visible devices at startup.

Summary

• In Python, GPU control usually means restricting device visibility rather than setting an abstract count.
• 'CUDA_VISIBLE_DEVICES is the standard mechanism and should be set before framework initialization.'
• TensorFlow can also restrict visible devices in code if it happens early enough.
• PyTorch usually inherits visibility from the environment and should verify it at runtime.
• Multi-GPU training requires explicit framework logic in addition to visible hardware.