Estas funciones deberían ayudar:

 >>> import torch >>> torch.cuda.is_available() True >>> torch.cuda.device_count() 1 >>> torch.cuda.current_device() 0 >>> torch.cuda.device(0) <torch.cuda.device at 0x7efce0b03be0> >>> torch.cuda.get_device_name(0) 'GeForce GTX 950M'

Esto nos dice:

• CUDA está disponible y puede ser utilizado por un dispositivo.
• Device 0 se refiere a la GPU GeForce GTX 950M , y PyTorch lo elige actualmente.

Como no se ha propuesto aquí, estoy agregando un método usando torch.device , ya que es bastante útil, también cuando se inicializan los tensores en el device correcto.

 # setting device on GPU if available, else CPU device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') print('Using device:', device) print() #Additional Info when using cuda if device.type == 'cuda': print(torch.cuda.get_device_name(0)) print('Memory Usage:') print('Allocated:', round(torch.cuda.memory_allocated(0)/1024**3,1), 'GB') print('Cached: ', round(torch.cuda.memory_reserved(0)/1024**3,1), 'GB')

Editar: torch.cuda.memory_cached ha sido renombrado a torch.cuda.memory_reserved . Así que usa memory_cached para versiones anteriores.

Producción:

 Using device: cuda Tesla K80 Memory Usage: Allocated: 0.3 GB Cached: 0.6 GB

Como se mencionó anteriormente, usando el device es posible :

• Para mover tensores al device respectivo:

   torch.rand(10).to(device)

• Para crear un tensor directamente en el device :

   torch.rand(10, device=device)

Lo que hace que cambiar entre CPU y GPU sea cómodo sin cambiar el código real.

Editar:

Como ha habido algunas preguntas y confusión sobre la memoria almacenada en caché y asignada , estoy agregando información adicional al respecto:

• torch.cuda.max_memory_cached(device=None)
  
  Devuelve la memoria de GPU máxima administrada por el asignador de almacenamiento en caché en bytes para un dispositivo determinado.
  

• torch.cuda.memory_allocated(device=None)
  
  Devuelve el uso de memoria GPU actual por tensores en bytes para un dispositivo determinado.

Puede entregar directamente un device como se especifica más arriba en la publicación o puede dejarlo como Ninguno y utilizará current_device() .

Nota adicional: ¡Las tarjetas gráficas antiguas con capacidad de cómputo Cuda 3.0 o inferior pueden estar visibles pero no pueden ser utilizadas por Pytorch!
¡Gracias a hekimgil por señalar esto! - "GPU0 GeForce GT 750M encontrada que tiene capacidad cuda 3.0. PyTorch ya no es compatible con esta GPU porque es demasiado antigua. La capacidad cuda mínima que admitimos es 3.5".

Después de comenzar a ejecutar el ciclo de entrenamiento, si desea verlo manualmente desde la terminal si su programa está utilizando los recursos de la GPU y en qué medida, simplemente puede usar el watch como en:

 $ watch -n 2 nvidia-smi

Esto actualizará continuamente las estadísticas de uso cada 2 segundos hasta que presione ctrl + c

Si necesita más control sobre más estadísticas de GPU que pueda necesitar, puede usar una versión más sofisticada de nvidia-smi con --query-gpu=... . A continuación se muestra una ilustración simple de esto:

 $ watch -n 3 nvidia-smi --query-gpu=index,gpu_name,memory.total,memory.used,memory.free,temperature.gpu,pstate,utilization.gpu,utilization.memory --format=csv

lo que generaría las estadísticas algo así como:

Nota : No debe haber ningún espacio entre los nombres de consulta separados por comas en --query-gpu=... . De lo contrario, esos valores se ignorarán y no se devolverán estadísticas.

Además, puede verificar si su instalación de PyTorch detecta su instalación de CUDA correctamente haciendo lo siguiente:

 In [13]: import torch In [14]: torch.cuda.is_available() Out[14]: True

El estado True significa que PyTorch está configurado correctamente y está usando la GPU, aunque debe mover/colocar los tensores con las declaraciones necesarias en su código.

Si desea hacer esto dentro del código de Python, consulte este módulo:

https://github.com/jonsafari/nvidia-ml-py o en pypi aquí: https://pypi.python.org/pypi/nvidia-ml-py/