Fonctions torch.save et torch.load de PyTorch
Manuel de référence Pytorch torch
torch.saveettorch.loadsont des fonctions de PyTorch pour sérialiser (sauvegarder) et désérialiser (charger) des tenseurs, des modèles et d'autres objets Python.
Ces fonctions sont essentielles dans des scénarios tels que la sauvegarde de modèles entraînés et la sauvegarde de points de contrôle pour reprendre l'entraînement.
Définition de fonction
torch.save(obj, f, pickle_module, pickle_protocol) torch.load(f, map_location, pickle_module, weights_only)
Paramètres de torch.save:
obj: L'objet à sauvegarder, qui peut être un tenseur, un modèle, un dictionnaire, une liste, etc.f: Chemin de fichier (chaîne ou objet fichier).pickle_module(Facultatif) : Le module utilisé pour la sérialisation.pickle_protocol(Facultatif) : Version du protocole de sérialisation.
Paramètres de torch.load:
f: Chemin de fichier (chaîne ou objet fichier).map_location(Facultatif) : Spécifie comment mapper le stockage vers différents appareils.pickle_module(Facultatif) : Le module utilisé pour la désérialisation.weights_only(bool, facultatif) : Indique s'il faut charger uniquement les poids et non les objets Python.
Exemples d'utilisation
Exemple 1 : Enregistrement et chargement de tenseurs
Exemple
import torch
# Créer quelques tenseurs
x = torch.tensor([1, 2, 3, 4, 5])
y = torch.randn(3, 4)
# Enregistrer les tenseurs dans un fichier
torch.save({'x': x, 'y': y}, 'tensors.pth')
# Charger les tenseurs depuis un fichier
loaded = torch.load('tensors.pth')
print("Données chargées :", loaded)
print("x:", loaded['x'])
print("y:", loaded['y'])
# Créer quelques tenseurs
x = torch.tensor([1, 2, 3, 4, 5])
y = torch.randn(3, 4)
# Enregistrer les tenseurs dans un fichier
torch.save({'x': x, 'y': y}, 'tensors.pth')
# Charger les tenseurs depuis un fichier
loaded = torch.load('tensors.pth')
print("Données chargées :", loaded)
print("x:", loaded['x'])
print("y:", loaded['y'])
La sortie est :
加载的数据: {'x': tensor([1, 2, 3, 4, 5]), 'y': tensor([[-0.3042, -0.9077, -1.0826, 0.9333],
[ 0.0551, 0.6728, 0.5942, -0.1522],
[-0.3744, 0.9239, -0.2104, -0.5239]])}
x: tensor([1, 2, 3, 4, 5])
y: tensor([[-0.3042, -0.9077, -1.0826, 0.9333],
[ 0.0551, 0.6728, 0.5942, -0.1522],
[-0.3744, 0.9239, -0.2104, -0.5239]])
Exemple 2 : Enregistrement et chargement de modèles
Exemple
import torch
import torch.nn as nn
# Définir un réseau de neurones simple
class SimpleNet(nn.Module):
def __init__(self):
super(SimpleNet, self).__init__()
self.fc1 = nn.Linear(10, 5)
self.fc2 = nn.Linear(5, 2)
def forward(self, x):
x = torch.relu(self.fc1(x))
x = self.fc2(x)
return x
# Créer une instance de modèle
model = SimpleNet()
# Enregistrer le modèle (enregistrer le modèle entier)
torch.save(model, 'model.pth')
# Charger le modèle
loaded_model = torch.load('model.pth')
print("Modèle enregistré et chargé")
print(loaded_model)
import torch.nn as nn
# Définir un réseau de neurones simple
class SimpleNet(nn.Module):
def __init__(self):
super(SimpleNet, self).__init__()
self.fc1 = nn.Linear(10, 5)
self.fc2 = nn.Linear(5, 2)
def forward(self, x):
x = torch.relu(self.fc1(x))
x = self.fc2(x)
return x
# Créer une instance de modèle
model = SimpleNet()
# Enregistrer le modèle (enregistrer le modèle entier)
torch.save(model, 'model.pth')
# Charger le modèle
loaded_model = torch.load('model.pth')
print("Modèle enregistré et chargé")
print(loaded_model)
Exemple 3 : Sauvegarde des paramètres du modèle uniquement (méthode recommandée)
Exemple
import torch
import torch.nn as nn
class SimpleNet(nn.Module):
def __init__(self):
super(SimpleNet, self).__init__()
self.fc1 = nn.Linear(10, 5)
self.fc2 = nn.Linear(5, 2)
def forward(self, x):
x = torch.relu(self.fc1(x))
x = self.fc2(x)
return x
model = SimpleNet()
# Sauvegarder uniquement les paramètres du modèle (méthode recommandée)
torch.save(model.state_dict(), 'model_weights.pth')
# Créer un nouveau modèle et charger les paramètres
new_model = SimpleNet()
new_model.load_state_dict(torch.load('model_weights.pth'))
print("Paramètres du modèle sauvegardés et chargés")
print(new_model.state_dict().keys())
import torch.nn as nn
class SimpleNet(nn.Module):
def __init__(self):
super(SimpleNet, self).__init__()
self.fc1 = nn.Linear(10, 5)
self.fc2 = nn.Linear(5, 2)
def forward(self, x):
x = torch.relu(self.fc1(x))
x = self.fc2(x)
return x
model = SimpleNet()
# Sauvegarder uniquement les paramètres du modèle (méthode recommandée)
torch.save(model.state_dict(), 'model_weights.pth')
# Créer un nouveau modèle et charger les paramètres
new_model = SimpleNet()
new_model.load_state_dict(torch.load('model_weights.pth'))
print("Paramètres du modèle sauvegardés et chargés")
print(new_model.state_dict().keys())
La sortie est :
模型参数已保存并加载 odict_keys(['fc1.weight', 'fc1.bias', 'fc2.weight', 'fc2.bias'])It is recommended to save only the model parameters (
state_dict), rather than saving the entire model. This allows parameters to be reused between models with different architectures.Example 4: Transferring Between CPU and GPU
Example
import torch
# Assume the model was saved on GPU
if torch.cuda.is_available():
x = torch.randn(2, 3, device='cuda')
torch.save(x, 'tensor_gpu.pth')
# Load on CPU
x_cpu = torch.load('tensor_gpu.pth', map_location='cpu')
print("Loaded to CPU:", x_cpu.device)
Using
map_locationthe parameter, you can load tensors onto different devices.
Notes
- Save files usually use
.pthou.ptextension. - It is recommended to save only the model parameters (
state_dict), rather than saving the entire model. - When loading, security and version compatibility should be considered.