pytorch save model like keras

Question

I wanna saving the model I trained in pytorch. And I want to save the model and the weight after training.

In keras I used

from tensorflow import keras

model_json = model.to_json()
with open("bpnn_model.json", "w") as json_file:
    json_file.write(model_json)

# saving weight
model.save_weights("bpnn_model_weights.h5")

to saving weight and model shape. And loading model I use

from tensorflow import keras

with open('bpnn_model.json', 'r') as file:
    model_json = file.read()

trained_model = keras.models.model_from_json(model_json)
trained_model.load_weights('bpnn_model_weights.h5')

So that I can easily rebuild the model without using the saving_model_code into my new code

In pytorch I tried torch.save() / model.state_dict() to save model. But all I get is Can't get attribute 'net' on <module '__main__' from D:\\... I asked chat-gpt and he told me need to build the model before loading the code. Or need to import the net in from building model code likefrom your_model_file import Net

What i want is when I open a new code, I just need to load the model and weight, and maybe setting the activation and optimizer. And I can retrain or fit the model, don't need to care about the model shape.

PS : My English are not very good , If i write something wrong or hard to understand, I apologize abuot that.

Answer 1

you can try torch.jit

import torch
from YourModel import yourModel

model = yourModel()
...
#training
...
#save model
model.eval()
jit_model = torch.jit.trace(model, torch.randn(N, C, H, W))
torch.jit.save(jit_model, 'your_path.pth')

after saving the jit model, you can load it like this:

import torch
...
#load model
model = torch.jit.load("./your_path.pth")