Errors removing __init__.py when using Arcface?

Hello,

I have a test.py file that uses “class ArcMarginModel(torch.nn.Module)” from an insightface_models.py in a folder named “models” in my project. This is imported in an init.py file as follows:

from .insightface_models import ArcMarginModel

When I remove the init.py I get the following error (even if I move everything from insightface_models.py into my test.py):

Message=Can’t get attribute ‘ArcMarginModel’ on <module ‘models’ (namespace)>
Source=B:\test_project\models\test.py

My code is as follows:

    arcface_checkpoint = torch.load(arcface_model, map_location=device)
    arcface = arcface_checkpoint['model'].module
    arcface = arcface.to(device)
    self.arcface.eval()

How can I make this work without the init.py? Why am I getting this error?

Thanks!

I guess you are trying to directly store and load the model via:

model = ArcMarginModel()
torch.save(model, ...)
model = torch.load(path)

As described in this tutorial you would have to make sure to restore the exact file structure and thus using the state_dict instead is the recommended way:

This save/load process uses the most intuitive syntax and involves the least amount of code. Saving a model in this way will save the entire module using Python’s pickle module. The disadvantage of this approach is that the serialized data is bound to the specific classes and the exact directory structure used when the model is saved. The reason for this is because pickle does not save the model class itself. Rather, it saves a path to the file containing the class, which is used during load time. Because of this, your code can break in various ways when used in other projects or after refactors.