Solución de memoria insuficiente CUDA

El enlace original es https://www.zhihu.com/question/274635237 , encontré CUDA sin memoria durante la prueba de detección de texto, y usé torch.no_grad () y torch.cuda.empty_cache () para reducir el uso de memoria de video El efecto es obvio.

Generalmente, la memoria de video ocupada se divide en tres partes:

1. La memoria de video ocupada por los propios parámetros del modelo de red.
2. Las variables intermedias o los parámetros generados durante el cálculo del modelo (incluido el avance / retroceso / optimizador) también ocupan la memoria de video.
3. El marco de programación en sí tiene algunos gastos generales adicionales.

Cambiar la estructura de la red

1. Reducir batch_size ...

2. Sacrifique la velocidad de cálculo para reducir el uso de memoria, divida el cálculo en dos mitades, primero calcule los resultados de la mitad del modelo, guarde los resultados intermedios y luego calcule la segunda mitad del modelo

   # 输入
   input = torch.rand(1, 10)
   # 假设我们有一个非常深的网络
   layers = [nn.Linear(10, 10) for _ in range(1000)]
   model = nn.Sequential(*layers)
   output = model(input)
   
   ### 可进行如下更改
   # 首先设置输入的input=>requires_grad=True
   # 如果不设置可能会导致得到的gradient为0
   
   input = torch.rand(1, 10, requires_grad=True)
   layers = [nn.Linear(10, 10) for _ in range(1000)]
   
   
   # 定义要计算的层函数，可以看到我们定义了两个
   # 一个计算前500个层，另一个计算后500个层
   
   def run_first_half(*args):
       x = args[0]
       for layer in layers[:500]:
           x = layer(x)
       return x
   
   def run_second_half(*args):
       x = args[0]
       for layer in layers[500:-1]:
           x = layer(x)
       return x
   
   # 我们引入新加的checkpoint
   from torch.utils.checkpoint import checkpoint
   
   x = checkpoint(run_first_half, input)
   x = checkpoint(run_second_half, x)
   # 最后一层单独调出来执行
   x = layers[-1](x)
   x.sum.backward()  # 这样就可以了
   

3. Utilice la agrupación para reducir el tamaño del mapa de características

4. Reducir el uso de capas completamente conectadas

No modifique la estructura de la red

1. Use operaciones inplace tanto como sea posible. Por ejemplo, relu puede usar inplace = True. Una forma simple de usarlo es la siguiente:

   def inplace_relu(m):
       classname = m.__class__.__name__
       if classname.find('Relu') != -1:
           m.inplace=True
   model.apply(inplace_relu)
   

2. Eliminar pérdida al final de cada ciclo, lo que puede ahorrar muy poca memoria de video, pero es mejor que nada

3. Con el cálculo mixto de precisión float16, puede ahorrar casi el 50% de la memoria de video, pero tenga cuidado con algunas operaciones inseguras, como media y suma, desbordamiento fp16

4. Que no requieren el bp Forward, como validación, uso de prueba torch.no_grad(), la atención model.eval()no es igual atorch.no_grad()

   • model.eval()Se notificará a todas las capas que está en modo de evaluación, de modo que las capas de abandono o de norma por lotes funcionarán en el modo de evaluación en lugar del modo de entrenamiento.
   • torch.no_grad()Afectar autogrado y deshabilitarlo. Reducirá el uso de memoria y acelerará los cálculos, y no se puede propagar hacia atrás (no es necesario en los scripts de evaluación).

5. torch.cuda.empty_cache () es una versión avanzada de del.

6. El uso de la transformación del optimizador, teóricamente sgd <momentum <adam, se puede ver en la fórmula de cálculo que hay variables intermedias adicionales

7. Convolución en profundidad

8. No cargue los datos de una vez, sino léalos parcialmente, para que básicamente no haya problemas de escasez de memoria.