Can't find class when load model

1

Hello!
I created, trained and saved my model MyModel in Jupiter as 'mymodel.pt'. But then I want to load it in another application. To do this, I created the same class in the file where I load it:

class MyModel(nn.Module):
    def __init__(self):
        ...
    def forward(self, x):
        ...

This file is the file with the predictor class, which is run through uvicorn. And when I run torch.load(path_to_pt), I see an error that there is no class MyModel:

File "/usr/local/lib/python3.8/dist-packages/torch/serialization.py", line 851, in _load
    result = unpickler.load()
AttributeError: Can't get attribute 'MyModel' on <module '__main__' from '/usr/local/bin/uvicorn'>

I understand that the MyModel class is called differently:

>>> print(MyModel)
# <class 'classification.my_model.MyModel'>

And I understand that this shouldn’t be a problem when working through env, but is there any way I can solve this problem without going to env?

2

It seems you’ve stored the model directly (not its staste_dict), so I think you would need to make sure the folder structure etc. on the inference application is the same as was used during training.
The recommended way would be to store the state_dict of models instead, recreate the model object in another script, and load its state_dict later, which would avoid these pickling issues.

1 Like