Метод index
Метод index класса device возвращает целочисленный индекс устройства.
Для CPU этот индекс всегда равен 0, а для GPU индекс соответствует номеру
видеокарты в системе (начиная с 0). Этот метод полезен при работе с
многопоточными вычислениями и распределённым обучением.
Синтаксис
device.index
Пример
Давайте получим индекс устройства для CPU:
import torch
device = torch.device('cpu')
idx = device.index
print(idx)
Результат выполнения кода:
0
Пример
Получим индекс для первого GPU в системе:
import torch
device = torch.device('cuda:0')
idx = device.index
print(idx)
Результат выполнения кода:
0
Пример
Проверим, что метод возвращает None, если индекс не указан явно:
import torch
device = torch.device('cuda')
idx = device.index
print(idx)
Результат выполнения кода:
None
Пример
Практический пример использования индекса для выбора устройства:
import torch
devices = [torch.device('cuda:0'), torch.device('cuda:1')]
for device in devices:
if device.index is not None:
print(f'Using GPU {device.index}')
else:
print('Using default GPU')
Результат выполнения кода:
"Using GPU 0"
"Using GPU 1"