Solution de mémoire insuffisante CUDA

Le lien d'origine est https://www.zhihu.com/question/274635237 , j'ai rencontré CUDA à court de mémoire pendant le test de détection de texte et utilisé torch.no_grad () et torch.cuda.empty_cache () pour réduire l'utilisation de mémoire vidéo. L'effet est évident.

Généralement, la mémoire vidéo occupée est divisée en trois parties:

1. La mémoire vidéo occupée par les propres paramètres du modèle de réseau.
2. Les variables ou paramètres intermédiaires générés lors du calcul du modèle (y compris avant / arrière / optimiseur) occupent également la mémoire vidéo.
3. Le cadre de programmation lui-même a des frais généraux supplémentaires.

Changer la structure du réseau

1. Réduisez batch_size ...

2. Sacrifiez la vitesse de calcul pour réduire l'utilisation de la mémoire, divisez le calcul en deux moitiés, calculez d'abord les résultats de la moitié du modèle, enregistrez les résultats intermédiaires, puis calculez la dernière moitié du modèle

   # 输入
   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. Utilisez la mise en commun pour réduire la taille de la carte des caractéristiques

4. Réduisez l'utilisation de couches entièrement connectées

Ne modifiez pas la structure du réseau

1. Utilisez autant que possible les opérations inplace. Par exemple, relu peut utiliser inplace = True. Voici une façon simple de l'utiliser:

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

2. Supprimer la perte à la fin de chaque cycle, ce qui peut économiser très peu de mémoire vidéo, mais c'est mieux que rien

3. En utilisant le calcul mixte de précision float16, vous pouvez économiser près de 50% de la mémoire vidéo, mais faites attention à certaines opérations dangereuses, telles que la moyenne et la somme, le dépassement fp16

4. Cela ne nécessite pas le bp avant, comme la validation, l'utilisation du test torch.no_grad(), l'attention model.eval()n'est pas égale àtorch.no_grad()

   • model.eval()Toutes les couches seront informées que vous êtes en mode évaluation, de sorte que la norme de lot ou les couches d'exclusion fonctionnent en mode évaluation au lieu du mode entraînement.
   • torch.no_grad()Affectez autograd et désactivez-le. Cela réduira l'utilisation de la mémoire et accélérera les calculs, et ne peut pas être rétropropagé (non nécessaire dans les scripts eval).

5. torch.cuda.empty_cache () est une version avancée de del.

6. L'utilisation de la transformation de l'optimiseur, théoriquement sgd <momentum <adam, on peut voir à partir de la formule de calcul qu'il existe des variables intermédiaires supplémentaires

7. Convolution en profondeur

8. Ne chargez pas les données en une seule fois, mais lisez-les partiellement, de sorte qu'il n'y ait pratiquement aucun problème de manque de mémoire