Почему параметр устройства torch.cuda.device, вероятно, больше, чем torch.cuda.device_count ()?

Question

Это код:

torch.cuda.device_count()
# output: 1
d2 = torch.cuda.device(2)
# d2.idx = 2

Это для модели, работающей на другом устройстве?

Answer 1

torch.cuda.device_count() возвращает количество доступных графических процессоров, тогда как torch.cuda.device(device) - это context-manager. Это облегчает правильную обработку ресурсов, то есть автоматическую настройку и освобождение ресурсов после использования. Таким образом, аргумент, который вы передаете этой функции, не означает, что вы переключаетесь / настраиваетесь на это устройство. Чтобы настроить текущее устройство, вы должны использовать torch.cuda.set_device(device). Хотя, согласно официальной документации здесь ,

Использование этой функции не рекомендуется в пользу устройства. В большинстве В таких случаях лучше использовать переменную среды CUDA_VISIBLE_DEVICES.

Чтобы узнать больше о torch.cuda.device () и о том, как он работает, пройдите это официальное обсуждение pytorch.