Ventana interactiva de Python de VSCode: solucionar problemas de soporte de GPU Nvidia para TensorFlow / PyTorch

Solo quiero darte la solución a mi problema, ya que me llevó bastante tiempo resolverlo y, como siempre, la solución fue muy simple.

También aborda este problema de GitHub que se cerró sin respuesta: https://github.com/microsoft/vscode-jupyter/issues/1697

Problema

Configuré correctamente la instalación local del controlador nvidia y tensorflow.

Ejecutar este fragmento de código en la terminal de Python o en un archivo .py en una terminal (interna o externa de vscode) resultó en la siguiente salida:

import tensorflow as tf
gpus = tf.config.list_physical_devices('GPU')
print(gpus)

Resultado: [PhysicalDevice(name=’/physicaldevice:GPU:0′, device_type=’GPU’)]

Pero dentro de la ventana interactiva de Python (cuaderno Jupyter), al hacer clic en -> Ejecutar celda, solo recibí una lista vacía [].

¿Tienes el mismo problema?

1. Observa la salida del cuaderno Jupyter:
   • Vscode -> Terminal -> Sección de salida -> Cambia en la parte superior derecha a Jupyter y observa el mensaje de error, por ejemplo, una biblioteca no se encontró.
2. Observa la variable de entorno LD_LIBRARY_PATH, que se utiliza para establecer las rutas a las bibliotecas C y que normalmente se establece/amplía en el archivo .bashrc durante la instalación de nvidia cuDNN.

En tu terminal:

echo $LD_LIBRARY_PATH

En tu celda interactiva de ipython/jupyter vscode:

import os
os.environ["LD_LIBRARY_PATH"]

Si recibes un error de clave, eso significa que la variable no está establecida.

Deberían tener la misma salida, por ejemplo: /usr/lib/cuda/include:/usr/lib/cuda/lib64:

1. Si difieren o obtienes un KeyError, entonces puedo ayudarte.
   Difieren porque el cuaderno Jupyter no se inicia como una shell de inicio de sesión y, por lo tanto, no carga el bashrc y tampoco tiene las variables exportadas.

Solución

La solución es crear un archivo .env en el directorio raíz de Python de tu proyecto e incluir este contenido:

LD_LIBRARY_PATH=/usr/lib/cuda/include:/usr/lib/cuda/lib64:

Recarga la ventana de VSCode y voilà, funcionará sin errores.

Encontré la solución en estos dos enlaces:
– https://github.com/jupyter/notebook/issues/5871 -> Solución para ejecutar el cuaderno Jupyter externamente (no dentro del cuaderno Jupyter)
– https://github.com/microsoft/vscode-jupyter/issues/1467#issuecomment-767914345