¿Cómo mover un tensor de antorcha de CPU a GPU y viceversa en Python?

En este artículo, veremos cómo mover un tensor de CPU a GPU y de GPU a CPU en Python.

¿Por qué necesitamos mover el tensor?

Esto se hace por las siguientes razones: 

• Cuando entrenamos grandes redes neuronales, necesitamos usar nuestra GPU para un entrenamiento más rápido. PyTorch espera que los datos se transfieran de la CPU a la GPU. Inicialmente, todos los datos están en la CPU. 
• Después de realizar todos los procesos relacionados con el entrenamiento, el tensor de salida también se produce en la GPU. A menudo, los resultados de nuestras redes neuronales necesitan preprocesamiento. La mayoría de las bibliotecas de preprocesamiento no tienen soporte para tensores y esperan una array NumPy. NumPy no almacena datos en la GPU, por lo que espera que los datos estén en la CPU. 

Ahora que sabemos por qué necesitamos hacer estas operaciones, veamos cómo podemos hacerlo en el tensor de PyTorch.

CPU a GPU

Para mover nuestros tensores de la CPU a la GPU, usamos cualquiera de estos comandos:

Tensor.cuda() 

Tensor.to("cuda")

Ejemplo: 

GPU a CPU 

Ahora, para mover nuestros tensores de GPU a CPU, hay dos condiciones: 

1. Tensor con required_grad = False, o
2. Tensor con required_grad = True

Ejemplo 1: Si required_grad = False , simplemente puede hacerlo como:

Tensor.cpu()

Ejemplo 2: Si required_grad = True, entonces necesita usar: 

Tensor.detach().cpu()