Проверка использования GPU в PyTorch

Зачастую при работе с библиотекой PyTorch для машинного обучения возникает необходимость убедиться в том, что библиотека использует возможности графического процессора (GPU). Это может быть важно, так как обучение моделей на GPU обычно происходит значительно быстрее, чем на центральном процессоре (CPU).

Обычно для проверки активности GPU используется консольная команда nvidia-smi. Однако, возможно, есть желание проверить использование GPU прямо из скрипта на Python. Например, для автоматизации процесса проверки или для внедрения такой проверки в код проекта.

В PyTorch есть встроенные средства для проверки того, используется ли GPU. Для этого можно воспользоваться следующим кодом:

import torch

print(torch.cuda.is_available())

Если этот код возвращает True, это означает, что PyTorch успешно обнаруживает GPU и может использовать его для обучения моделей. Если возвращается False, значит, GPU либо отсутствует, либо PyTorch не может его обнаружить.

Также можно проверить, на каком устройстве находится тензор. Для этого можно использовать метод .device у тензора:

x = torch.Tensor([1, 2, 3])
print(x.device)

Возвращаемое значение будет указывать на устройство, на котором находится тензор: cpu для центрального процессора и cuda:0 для GPU.