Ce code renverra de la mémoire GPU libre en mégaoctets pour chaque GPU :
import subprocess as sp
import os
def get_gpu_memory():
_output_to_list = lambda x: x.decode('ascii').split('\n')[:-1]
ACCEPTABLE_AVAILABLE_MEMORY = 1024
COMMAND = "nvidia-smi --query-gpu=memory.free --format=csv"
memory_free_info = _output_to_list(sp.check_output(COMMAND.split()))[1:]
memory_free_values = [int(x.split()[0]) for i, x in enumerate(memory_free_info)]
print(memory_free_values)
return memory_free_values
get_gpu_memory()
Cette réponse repose sur l'installation de nvidia-smi (ce qui est presque toujours le cas pour les GPU Nvidia) et est donc limitée aux GPU NVidia.
En fait, j'ai trouvé une réponse dans cette vieille question. Pour apporter un avantage supplémentaire aux lecteurs, j'ai testé le programme mentionné
import nvidia_smi
nvidia_smi.nvmlInit()
handle = nvidia_smi.nvmlDeviceGetHandleByIndex(0)
# card id 0 hardcoded here, there is also a call to get all available card ids, so we could iterate
info = nvidia_smi.nvmlDeviceGetMemoryInfo(handle)
print("Total memory:", info.total)
print("Free memory:", info.free)
print("Used memory:", info.used)
nvidia_smi.nvmlShutdown()
sur colab avec le résultat suivant :
Total memory: 17071734784
Free memory: 17071734784
Used memory: 0
Le GPU que j'avais là-bas était un Tesla P100, comme on peut le voir en exécutant
!nvidia-smi
et observer la sortie
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 440.44 Driver Version: 418.67 CUDA Version: 10.1 |
|-------------------------------+----------------------+----------------------+
| GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |
|===============================+======================+======================|
| 0 Tesla P100-PCIE... Off | 00000000:00:04.0 Off | 0 |
| N/A 32C P0 26W / 250W | 0MiB / 16280MiB | 0% Default |
+-------------------------------+----------------------+----------------------+
+-----------------------------------------------------------------------------+
| Processes: GPU Memory |
| GPU PID Type Process name Usage |
|=============================================================================|
| No running processes found |
+-----------------------------------------------------------------------------+
