Функция mps.is_available

Функция mps.is_available используется для проверки доступности бэкенда MPS (Metal Performance Shaders) на устройствах Apple с macOS. Функция возвращает булево значение: True, если MPS доступен и может использоваться для вычислений, и False в противном случае. MPS позволяет ускорить обучение и вывод нейронных сетей на графических процессорах Apple Silicon (M1, M2, M3 и новее). Функция не требует передачи параметров.

Синтаксис

torch.mps.is_available()

Пример с проверкой доступности

Давайте проверим, доступен ли MPS на текущем устройстве:

import torch

res = torch.mps.is_available()
print(res)

Результат выполнения кода на устройстве с Apple Silicon:

True

Пример с условным использованием MPS

Давайте создадим тензор на устройстве MPS, если оно доступно, иначе используем CPU:

import torch

device = torch.device(
	"mps" if torch.mps.is_available()
	else "cpu"
)
t = torch.tensor([1, 2, 3, 4, 5]).to(device)
print(t)

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

tensor([1, 2, 3, 4, 5], device='mps:0')

Пример с синхронизацией MPS

При использовании MPS может потребоваться явная синхронизация для корректного выполнения операций:

import torch

if torch.mps.is_available():
	device = torch.device("mps")
	t1 = torch.tensor([1, 2, 3]).to(device)
	t2 = torch.tensor([4, 5, 6]).to(device)
	t_sum = t1 + t2
	torch.mps.synchronize()
	print(t_sum)

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

tensor([5, 7, 9], device='mps:0')

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

функцию mps.is_built,
которая проверяет, была ли собрана поддержка MPS в PyTorch
функцию mps.synchronize,
которая выполняет синхронизацию для бэкенда MPS
функцию cuda.is_available,
которая проверяет доступность бэкенда CUDA
функцию cuda.device_count,
которая возвращает количество доступных устройств CUDA