Pytorch: Demande de fonctionnalité : load_state_dict doit prendre des noms de fichiers

Créé le 31 mai 2017 · 3Commentaires · Source: pytorch/pytorch

Dans les situations de forte pression de mémoire, ce qui suit est un phénomène courant :

créer un modèle
lire state_dict à partir du fichier de point de contrôle (charge sur le GPU)
model.load_state_dict(s)

En raison de la pression de la mémoire, une solution de contournement courante consiste à faire d'abord :

s = torch.load('my_file.pt', map_location=lambda storage, loc: storage)

Et puis chargez s dans model .

C'est un scénario très courant que nous devrions être en mesure d'éviter, et ce scénario peut comporter quelques pièges : ce qui se passe sur les modèles mi-GPU, mi-CPU, ce qui se passe sur les modèles multi-GPU...

si load_state_dict a pris un nom de fichier directement, il peut supprimer ses stockages de paramètres existants et les définir sur le nouveau à la volée, ne nécessitant ainsi aucune mémoire supplémentaire.

feature nn triaged

Source

soumith

Commentaire le plus utile

Si load_state_dict prend un nom de fichier, nous devrions également autoriser le map_location . Une situation courante pour moi est de sauvegarder un point de contrôle sur la machine du cluster, puis de le charger sur mon macbook (il faut donc charger les paramètres sur le processeur)

alykhantejani le 1 juin 2017

👍3

Tous les 3 commentaires

il en va de même pour les state_dicts de l'optimiseur. pour certains optimiseurs comme Adagrad, les points de contrôle sont volumineux et nous pouvons avoir la même situation de pression mémoire. les optimiseurs n'ont même pas de .cuda() , nous devons donc d'abord charger manuellement state_dict sur le CPU, puis copier manuellement les parties sur le GPU.

J'ai rencontré cela en aidant @aszlam aujourd'hui.

soumith le 31 mai 2017

alykhantejani le 1 juin 2017

👍3

Moi et @szagoruyko sommes fans du format HDF5 pour les modèles sérialisés, peut-être que si cela pouvait bien s'entendre avec cette proposition