Атрибут device
Атрибут device класса Tensor возвращает объект
torch.device, который указывает на устройство,
где размещён данный тензор. Это может быть центральный
процессор (CPU), графический процессор NVIDIA (CUDA)
или графический процессор Apple (MPS). Атрибут доступен
только для чтения и автоматически обновляется при
переносе тензора с помощью методов to,
cuda или cpu.
Синтаксис
torch.device = tensor.device
Пример
Давайте создадим тензор на CPU и получим его устройство:
import torch
t = torch.tensor([1, 2, 3, 4, 5])
device = t.device
print(device)
Результат выполнения кода:
device(type='cpu')
Пример
Теперь создадим тензор на GPU с индексом 0:
import torch
if torch.cuda.is_available():
t = torch.tensor([1, 2, 3, 4, 5], device='cuda:0')
print(t.device)
Результат выполнения кода:
device(type='cuda', index=0)
Пример
Атрибут device часто используется для переноса
других тензоров на то же устройство:
import torch
t1 = torch.tensor([1, 2, 3, 4, 5])
t2 = torch.tensor([10, 20, 30, 40, 50])
t2 = t2.to(t1.device)
print(t2.device)
Результат выполнения кода:
device(type='cpu')
Пример
Проверим, поддерживается ли устройство MPS (Apple Silicon):
import torch
if torch.backends.mps.is_available():
t = torch.tensor([1, 2, 3, 4, 5], device='mps')
print(t.device)
Результат выполнения кода:
device(type='mps')