Функция cuda.device_count

Функция cuda.device_count возвращает количество доступных GPU-устройств CUDA в системе. Функция не принимает никаких аргументов. Возвращаемое значение может быть нулевым, если в системе нет доступных GPU с поддержкой CUDA. Перед вызовом этой функции рекомендуется проверить доступность CUDA с помощью функции cuda.is_available.

Синтаксис

torch.cuda.device_count()

Пример

Давайте получим количество доступных GPU в системе и выведем результат:

import torch

device_count = torch.cuda.device_count()
print(device_count)

Результат выполнения кода на системе с двумя GPU:

2

Пример

Давайте проверим доступность CUDA перед подсчетом устройств:

import torch

if torch.cuda.is_available():
	device_count = torch.cuda.device_count()
	print(f"Found {device_count} GPU(s)")
else:
	print("CUDA is not available")

Результат выполнения кода на системе с одним GPU:

"Found 1 GPU(s)"

Пример

На практике количество устройств часто используется для распределения нагрузки между несколькими GPU. Рассмотрим пример, который выводит названия всех доступных устройств:

import torch

if torch.cuda.is_available():
	gpu_count = torch.cuda.device_count()
	for i in range(gpu_count):
		gpu_name = torch.cuda.get_device_name(i)
		print(f"GPU {i}: {gpu_name}")
else:
	print("No CUDA devices available")

Результат выполнения кода на системе с двумя GPU:

"GPU 0: NVIDIA GeForce RTX 3080"
"GPU 1: NVIDIA GeForce RTX 3080"

Смотрите также

функцию cuda.is_available,
которая проверяет доступность CUDA в системе
функцию cuda.current_device,
которая возвращает индекс текущего устройства
функцию cuda.get_device_name,
которая возвращает имя GPU по его индексу
функцию cuda.set_device,
которая устанавливает активное устройство по индексу