Pytorch: Funktionsanfrage: load_state_dict sollte Dateinamen übernehmen
In Situationen mit hohem Speicherdruck tritt Folgendes häufig auf:
- Modell erstellen
- state_dict aus Checkpoint-Datei lesen (lädt auf GPU)
- model.load_state_dict(s)
Aufgrund des Speicherdrucks besteht eine gängige Problemumgehung darin, zuerst Folgendes zu tun:
s = torch.load('my_file.pt', map_location=lambda storage, loc: storage)
Laden Sie dann s in model .
Dies ist ein sehr häufiges Szenario, das wir vermeiden sollten, und dieses Szenario könnte einige Fallstricke haben: Was passiert bei Teil-GPU-Teil-CPU-Modellen, was passiert bei Multi-GPU-Modellen ...
Wenn load_state_dict direkt einen Dateinamen verwendet, kann es seine vorhandenen Parameterspeicher löschen und sie im laufenden Betrieb auf den neuen setzen, wodurch kein zusätzlicher Speicher benötigt wird.
soumith
Alle 3 Kommentare
dasselbe gilt für den Optimierer state_dicts. Bei einigen Optimierern wie Adagrad sind die Prüfpunkte groß, und wir können die gleiche Speicherdrucksituation haben. Optimierer haben nicht einmal ein .cuda() , also müssen wir zuerst manuell state_dict auf die CPU laden und dann Teile manuell auf die GPU kopieren.
Ich bin heute darauf gestoßen, als ich @aszlam geholfen habe.
soumith
am 31. Mai 2017
Wenn load_state_dict einen Dateinamen verwendet, sollten wir auch den Parameter map_location zulassen. Eine häufige Situation für mich ist, einen Prüfpunkt auf einem Cluster-Computer zu speichern und ihn dann auf mein MacBook zu laden (daher müssen Parameter auf die CPU geladen werden).
alykhantejani
am 1. Juni 2017
Ich und @szagoruyko sind Fans des HDF5-Formats für serialisierte Modelle, vielleicht könnte es mit diesem Vorschlag gut zurechtkommen
vadimkantorov
am 6. Juni 2017
Verwandte Themen
rajarshd
·
3Kommentare
keskarnitish
·
3Kommentare
kuangliu
·
3Kommentare
mishraswapnil
·
3Kommentare
eliabruni
·
3Kommentare
Hilfreichster Kommentar
Wenn
load_state_dicteinen Dateinamen verwendet, sollten wir auch den Parametermap_locationzulassen. Eine häufige Situation für mich ist, einen Prüfpunkt auf einem Cluster-Computer zu speichern und ihn dann auf mein MacBook zu laden (daher müssen Parameter auf die CPU geladen werden).