Loading model from .pth is not supported yet

[SkafteNicki]

🐛 Bug

To Reproduce

import torch
from litmodels import save_model, load_model

model = torch.nn.Module()
save_model(model=model, name="skaftenicki/language-model/torch-model")
model_ = load_model(name="skaftenicki/language-model/torch-model")

generates

❯ python main.py 
Uploading files...: 100%|█████████████████| 1/1 [00:01<00:00,  1.73s/it]
Model uploaded successfully. Link to the model: 'https://lightning.ai/skaftenicki/language-model/models/torch-model'
Downloading default: 100%|█████████| 1.32k/1.32k [00:00<00:00, 1.71kB/s]
Traceback (most recent call last):          | 0.00/1.32k [00:00<?, ?B/s]
  File "/home/nsde/Documents/LitModels/main.py", line 6, in <module>
    model_ = load_model(name="skaftenicki/language-model/torch-model")
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/nsde/Documents/LitModels/src/litmodels/io/gateway.py", line 166, in load_model
    raise NotImplementedError(f"Loading model from {model_path.suffix} is not supported yet.")
NotImplementedError: Loading model from .pth is not supported yet.

Expected behavior

Code is basically taken from the examples. The problem seems to be that save_model function saves torch models to pth (which are not ScriptModule):

LitModels/src/litmodels/io/gateway.py

Lines 94 to 96 in e9075da

	elif _PYTORCH_AVAILABLE and isinstance(model, torch.nn.Module):
	path = os.path.join(staging_dir, f"{model.__class__.__name__}.pth")
	torch.save(model.state_dict(), path)

but not supported in the load_model function.

Environment

• PyTorch Version (e.g., 1.0):
• OS (e.g., Linux):
• How you installed PyTorch (conda, pip, source):
• Build command you used (if compiling from source):
• Python version:
• CUDA/cuDNN version:
• GPU models and configuration:
• Any other relevant information:

Additional context

[Borda]

Good call, we have just savig

LitModels/src/litmodels/io/gateway.py

Line 94 in e9075da

elif _PYTORCH_AVAILABLE and isinstance(model, torch.nn.Module):

As it saves the dict, only you need the class to restore it... in such case I would suggest to use pickle or mixin:

LitModels/src/litmodels/integrations/mixins.py

Line 124 in e9075da

class PyTorchRegistryMixin(ModelRegistryMixin):

Or would you have an idea how to address this limitation?

[SkafteNicki]

@Borda yeah i quickly realized that the problem is that Pytorch needs the class to be initialized first and then use load_state_dict. My main issue is that if you look at the example from the readme:

LitModels/README.md

Lines 65 to 81 in bbd28a1

	Save model:
	```python
	import torch
	from litmodels import save_model
	model = torch.nn.Module()
	save_model(model=model, name="your_org/your_team/torch-model")
	```
	Load model:
	```python
	from litmodels import load_model
	model_ = load_model(name="your_org/your_team/torch-model")
	```

this does not work right now. So update docs/readme should probably be updated. However, I do think if the save_model function can save and upload to a specific format, that the load_model function should be able to work with that format.

One proposal is to add optional arg to load_model:

def load_model(name: str, download_dir: str = ".", initialized_model: Optional[object] = None) -> Any:
    ...
    if model_path.suffix.lower() == ".pth":
        if initialized_model is not None and isinstance(initialized_model, torch.nn.Module):
            initialized_model.load_state_dict(torch.load(model_path))
            return initialized_model

[Borda]

One proposal is to add optional arg to load_model

Yes, that looks as a good solution for now, mind send a PR :)