¿Cómo sabes si un archivo de guardado de Pytorch contiene un modelo y/o solo los pesos?

Answer

1. 

   davy.ai

   July 30, 2023 at 4:13 pm

   Si estás viendo la extensión de archivo “.pth” siendo utilizada tanto para torch.load como para torch.load_state_dict, es probable que el archivo contenga un modelo guardado de PyTorch o simplemente el diccionario de estado de un modelo de PyTorch. Sin embargo, como mencionaste que model.eval() falla después de cargarlo, es posible que el archivo solo contenga el diccionario de estado y no el modelo completo.

   Para inspeccionar el contenido del archivo y saber si contiene el modelo completo o solo el diccionario de estado, puedes usar el siguiente código:

   import torch
   
   # Cargar el modelo o el diccionario de estado desde el archivo
   loaded_dict = torch.load(file_path)
   
   # Verificar si el objeto cargado es un modelo o solo un diccionario de estado
   if isinstance(loaded_dict, dict) and 'model' in loaded_dict:
       # Si el objeto es un diccionario y contiene la clave 'model', es el modelo completo
       model = loaded_dict['model']
       model.eval()
   else:
       # Si el objeto no es un diccionario o no contiene la clave 'model', es solo el diccionario de estado
       model = TuClaseDeModeloPersonalizada()
       model.load_state_dict(loaded_dict)
       model.eval()
   

   En este código, primero cargamos el objeto desde el archivo usando torch.load. Luego verificamos si el objeto cargado es un diccionario y si contiene la clave ‘model’. Si lo hace, sabemos que el archivo contiene el modelo completo de PyTorch, por lo que podemos extraer el modelo del diccionario usando la clave ‘model’ y llamar a model.eval() en él. Si no contiene la clave ‘model’, sabemos que es solo el diccionario de estado, por lo que podemos crear una instancia del modelo usando tu clase de modelo personalizada TuClaseDeModeloPersonalizada() y cargar el diccionario de estado usando model.load_state_dict(loaded_dict) antes de llamar a model.eval() en él.

   ¡Espero que esto te ayude!

Comments are closed.