Funciones torch.save y torch.load de PyTorch

Pytorch torch 参考手册Manual de referencia de torch de PyTorch

torch.saveytorch.loadSon funciones en PyTorch utilizadas para serializar (guardar) y deserializar (cargar) tensores, modelos y otros objetos de Python.

Estas funciones son esenciales en situaciones como guardar modelos entrenados o guardar puntos de control para reanudar el entrenamiento.

Definición de funciones

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

Parámetros de torch.save:

  • obj: El objeto a guardar, puede ser un tensor, modelo, diccionario, lista, etc.
  • f: La ruta del archivo (cadena de caracteres u objeto de archivo).
  • pickle_module(Opcional): El módulo utilizado para la serialización.
  • pickle_protocol(Opcional): Versión del protocolo de serialización.

Parámetros de torch.load:

  • f: La ruta del archivo (cadena de caracteres u objeto de archivo).
  • map_location(Opcional): Especifica cómo asignar el almacenamiento a diferentes dispositivos.
  • pickle_module(Opcional): El módulo utilizado para la deserialización.
  • weights_only(bool, opcional): Indica si solo se cargan los pesos sin cargar objetos de Python.

Ejemplos de uso

Ejemplo 1: Guardar y cargar tensores

Ejemplo

import torch

# Crear algunos tensores
x = torch.tensor([1, 2, 3, 4, 5])
y = torch.randn(3, 4)

# Guardar el tensor en un archivo
torch.save({'x': x, 'y': y}, 'tensors.pth')

# Cargar el tensor desde un archivo
loaded = torch.load('tensors.pth')

print("Datos cargados:", loaded)
print("x:", loaded['x'])
print("y:", loaded['y'])

El resultado de salida es:

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

Ejemplo 2: Guardar y cargar modelos

Ejemplo

import torch
import torch.nn as nn

# Definir una red neuronal 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

# Crear una instancia del modelo
model = SimpleNet()

# Guardar el modelo (guardar el modelo completo)
torch.save(model, 'model.pth')

# Cargar el modelo
loaded_model = torch.load('model.pth')

print("El modelo se ha guardado y cargado")
print(loaded_model)

Ejemplo 3: Guardar solo los parámetros del modelo (método recomendado)

Ejemplo

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

# Guardar solo los parámetros del modelo (método recomendado)
torch.save(model.state_dict(), 'model_weights.pth')

# Crear un nuevo modelo y cargar los parámetros
new_model = SimpleNet()
new_model.load_state_dict(torch.load('model_weights.pth'))

print("Los parámetros del modelo se han guardado y cargado")
print(new_model.state_dict().keys())

El resultado de salida es:

模型参数已保存并加载
odict_keys(['fc1.weight', 'fc1.bias', 'fc2.weight', 'fc2.bias'])



Se recomienda guardar solo los parámetros del modelo (state_dict), en lugar de guardar todo el modelo. Así se pueden reutilizar los parámetros entre modelos de diferentes arquitecturas.



Ejemplo 4: Migrar entre CPU y GPU





Ejemplo



import torch



# Supongamos que el modelo se guardó en GPU

if torch.cuda.is_available():

    x = torch.randn(2, 3, device='cuda')

    torch.save(x, 'tensor_gpu.pth')



    # Cargar en CPU

    x_cpu = torch.load('tensor_gpu.pth', map_location='cpu')

    print("Cargado en CPU:", x_cpu.device)







Usandomap_locationel parámetro, se pueden cargar tensores en diferentes dispositivos.





Notas



    

  • Al guardar archivos, normalmente se usa.pth o  .ptla extensión.
    • 

  • Se recomienda guardar solo los parámetros del modelo (state_dict), en lugar de guardar todo el modelo.
    • 

  • Al cargar, se debe prestar atención a la seguridad y la compatibilidad de versiones.
    • 






Pytorch torch 参考手册Manual de referencia de Pytorch torch