Guardar los pesos del modelo por separado en PyTorch.
Estoy utilizando PyTorch para entrenar un modelo de aprendizaje profundo. Me pregunto si es posible guardar por separado los pesos del modelo. Por ejemplo:
class my_model(nn.Module):
def __init__(self):
super(my_model, self).__init__()
self.bert = transformers.AutoModel.from_pretrained(BERT_PATH)
self.out = nn.Linear(768,1)
def forward(self, ids, mask, token_type):
x = self.bert(ids, mask, token_type)[1]
x = self.out(x)
return x
Tengo el modelo BERT como base y una capa lineal adicional en la parte superior. Después de entrenar este modelo, ¿puedo guardar los pesos del modelo BERT y esta capa lineal por separado?
davy.ai
Sí, es posible guardar por separado los pesos del modelo BERT y la capa lineal en PyTorch. Para hacer esto, puedes usar el método
state_dict()
del modelo para obtener un diccionario que contenga el estado de los parámetros y buffers del modelo. Luego, puedes guardar el diccionario de estado utilizando la funcióntorch.save()
.Para guardar por separado los pesos del modelo BERT y la capa lineal, puedes usar el siguiente código:
Este código crea una instancia de
mi_modelo
, lo entrena y luego obtiene los diccionarios de estado para el modelo BERT y la capa lineal utilizando el métodostate_dict()
. Los diccionarios de estado se guardan en archivos separados utilizando la funcióntorch.save()
.Para cargar los pesos guardados para el modelo BERT y la capa lineal, puedes utilizar la función
torch.load()
para cargar los diccionarios de estado, y luego utilizar el métodoload_state_dict()
de los respectivos componentes del modelo para actualizar sus parámetros y buffers:Este código crea una instancia de
mi_modelo
y luego carga los diccionarios de estado guardados para el modelo BERT y la capa lineal utilizandotorch.load()
. Luego se utiliza el métodoload_state_dict()
de los respectivos componentes del modelo para actualizar sus parámetros y buffers con los diccionarios de estado cargados.