PyTorch: вывод dtype из возможностей устройства, а не входных данных

Question

Я пытаюсь написать библиотеку для устройств c для PyTorch, и я наткнулся на проблему PyTorch, используя dtypes, которые не совместимы с моим вычислительным устройством:

import scipy.signal
import torch

raw_window = scipy.signal.windows.cosine(128)
print(raw_window.dtype)  # float64

device = torch.device("cpu")
window = torch.as_tensor(raw_window, device=device)
print(window.device)  # cpu
print(window.dtype)   # torch.float64

device = torch.device("cuda")
window = torch.as_tensor(raw_window, device=device)
print(window.device)  # cuda:0
print(window.dtype)   # torch.float64

Как вы Как видно из последней строки, torch назначает dtype torch.float64, хотя мое устройство CUDA не может обрабатывать значения с плавающей запятой двойной точности.

Есть ли способ заставить PyTorch использовать вместо этого наиболее подходящий тип dtype устройства входных данных dtype? Или я неправильно понимаю всю концепцию установки устройств и dtypes полностью?

Answer 1

Нет, как вы заметили, PyTorch выводит dtype только из input данных.

В вашем случае для numpy по умолчанию установлено значение np.float64 (независимо от системы и архитектуры) PyTorch выведет, что он аналогичен torch.float64, поэтому больше проблем с началом с numpy (и вы не можете установить другое значение по умолчанию dtype).

In pytorch Вы обычно go для torch.float32 (и это значение по умолчанию), в конечном итоге torch.float16 для смешанной точности (исключая квантование). Обычно такой точности достаточно (хотя и не уверен в вашем конкретном сценарии использования).

Следовательно, ваш лучший выбор - cast numpy массивов до np.float32, например:

raw_window = raw_window.astype(np.float32)

(или с помощью pytorch или, что еще лучше, создайте float32 с самого начала), который AFAIK должен быть доступен на всех / большинстве устройств, будь то CPU или GPU.

Для агрегирования CPU / GPU c подходить вам следует go для явного создания устройства , как показано здесь :

if not args.disable_cuda and torch.cuda.is_available():
    args.device = torch.device('cuda')
else:
    args.device = torch.device('cpu')

(расширить, если вам нужно явно указать cuda или аналогичный) и использовать его во время создания tensor.