PyTorch-Funktionen torch.save und torch.load

Pytorch torch 参考手册PyTorch torch Referenzhandbuch

torch.saveundtorch.loadsind Funktionen in PyTorch zum Serialisieren (Speichern) und Deserialisieren (Laden) von Tensoren, Modellen und anderen Python-Objekten.

Diese Funktionen sind in Szenarien wie dem Speichern trainierter Modelle und dem Speichern von Checkpoints zur Wiederaufnahme des Trainings unerlässlich.

Funktionsdefinition

torch.save(obj, f, pickle_module, pickle_protocol)
torch.load(f, map_location, pickle_module, weights_only)

torch.save-Parameter:

  • obj: Das zu speichernde Objekt, das ein Tensor, Modell, Dictionary, eine Liste usw. sein kann.
  • f: Dateipfad (Zeichenkette oder Dateiobjekt).
  • pickle_module(Optional): Das zur Serialisierung verwendete Modul.
  • pickle_protocol(Optional): Version des Serialisierungsprotokolls.

torch.load-Parameter:

  • f: Dateipfad (Zeichenkette oder Dateiobjekt).
  • map_location(Optional): Gibt an, wie Speicher auf verschiedene Geräte abgebildet wird.
  • pickle_module(Optional): Das zur Deserialisierung verwendete Modul.
  • weights_only(bool, optional): Ob nur Gewichte und keine Python-Objekte geladen werden sollen.

Anwendungsbeispiele

Beispiel 1: Speichern und Laden von Tensoren

Beispiel

import torch

# Einige Tensoren erstellen
x = torch.tensor([1, 2, 3, 4, 5])
y = torch.randn(3, 4)

# Tensoren in Datei speichern
torch.save({'x': x, 'y': y}, 'tensors.pth')

# Tensoren aus Datei laden
loaded = torch.load('tensors.pth')

print("Geladene Daten:", loaded)
print("x:", loaded['x'])
print("y:", loaded['y'])

Die Ausgabe ist:

加载的数据: {'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]])

Beispiel 2: Speichern und Laden von Modellen

Beispiel

import torch
import torch.nn as nn

# Ein einfaches neuronales Netzwerk definieren
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

# Eine Modellinstanz erstellen
model = SimpleNet()

# Das Modell speichern (das gesamte Modell speichern)
torch.save(model, 'model.pth')

# Modell laden
loaded_model = torch.load('model.pth')

print("Modell gespeichert und geladen")
print(loaded_model)

Beispiel 3: Nur Modellparameter speichern (Empfohlene Methode)

Beispiel

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()

# Nur Modellparameter speichern (empfohlene Methode)
torch.save(model.state_dict(), 'model_weights.pth')

# Ein neues Modell erstellen und Parameter laden
new_model = SimpleNet()
new_model.load_state_dict(torch.load('model_weights.pth'))

print("Modellparameter gespeichert und geladen")
print(new_model.state_dict().keys())

Die Ausgabe ist:

模型参数已保存并加载
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)







Usingmap_locationthe parameter, you can load tensors onto different devices.





Notes



    

  • Save files usually use.pth oder  .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.
    • 






Pytorch torch 参考手册Pytorch torch Reference Manual