Importing models saved with pytorch

[schallerala]

Using the latest release 0.2.7, I tried to import models that have been saved using pytorch like so mlem import ckpt/best.pt ckpt/best which produced the following error:

❌ Unexpected error: A load persistent id instruction was encountered,
but no persistent_load function was specified.
Please report it here: <https://github.com/iterative/mlem/issues>

From my understanding, MLEM doesn't support the new way of pytorch to pickle its model using the torch.save method (https://pytorch.org/docs/stable/generated/torch.save.html). Most likely, when we read their note:

The 1.6 release of PyTorch switched torch.save to use a new zipfile-based file format. torch.load still retains the ability to load files in the old format. If for any reason you want torch.save to use the old format, pass the kwarg _use_new_zipfile_serialization=False.

Additionally, reading at MLEM's source, found out it is "only" able to import pickled models when looking at the subclasses of the ExtImportHook (model-wise: PickleImportHook) (https://github.com/iterative/mlem/blob/HEAD/mlem/core/import_objects.py#L56-L67).

Therefore, I guess I better change the source code and use MLEM's save method, to persist models; but believe it could be a great new feature of the import command :)

Leaving it here to the pros! 👨🏼‍🎤

Best,

[aguschin]

Thanks for reporting this, @schallerala! We'll take a look when we'll have more time to improve import functionality :)

[mike0sv]

@schallerala can you please try installing from branch above and running mlem import ckpt/best.pt ckpt/best --type torch?

[schallerala]

@mike0sv I fail to do so because when running the command you mentioned, I get the following message:

❌ Extension 'torch' requires additional dependencies: pip install mlem[torch] or pip install torch

Even though I installed your branch like so: pip install "mlem[torch] @ git+https://github.com/iterative/mlem.git@feature/torch-import" (like described in the pip install documentation: example 7) and being able to import torch and for example print cuda's architecture list:

$ python
>>> import torch
>>> torch.cuda.get_arch_list()
['sm_37', 'sm_50', 'sm_60', 'sm_70']

-- and yes, sure of the environment. Have GPUs, tried inside poetry, a virtualenv and even install globally.

From the code and tests you submit in your PR, I guess it must be all good, but failed to run what you asked me to

[mike0sv]

can you add --tb like mlem --tb import ... and share the output?

[mike0sv]

and also try with API?

from mlem.api import import_objects; import_objects(..., type_="torch")

[schallerala]

Ok. So 2 problems:

First try with --tb

Seems I was also missing numpy as a dependency on a new environment.

After installing numpy and trying again

$ mlem --tb import ckpt/best.pt ckpt/best --type torch
⏳️ Importing object from ckpt/best.pt
Traceback (most recent call last):
  File "/path/.local/bin/mlem", line 8, in <module>
    sys.exit(app())
  File "/path/.local/lib/python3.9/site-packages/typer/main.py", line 328, in __call__
    raise e
  File "/path/.local/lib/python3.9/site-packages/typer/main.py", line 311, in __call__
    return get_command(self)(*args, **kwargs)
  File "/path/.local/lib/python3.9/site-packages/click/core.py", line 1130, in __call__
    return self.main(*args, **kwargs)
  File "/path/.local/lib/python3.9/site-packages/typer/core.py", line 778, in main
    return _main(
  File "/path/.local/lib/python3.9/site-packages/typer/core.py", line 216, in _main
    rv = self.invoke(ctx)
  File "/path/.local/lib/python3.9/site-packages/click/core.py", line 1657, in invoke
    return _process_result(sub_ctx.command.invoke(sub_ctx))
  File "/path/.local/lib/python3.9/site-packages/click/core.py", line 1404, in invoke
    return ctx.invoke(self.callback, **ctx.params)
  File "/path/.local/lib/python3.9/site-packages/click/core.py", line 760, in invoke
    return __callback(*args, **kwargs)
  File "/path/.local/lib/python3.9/site-packages/typer/main.py", line 683, in wrapper
    return callback(**use_params)  # type: ignore
  File "/path/.local/lib/python3.9/site-packages/click/decorators.py", line 26, in new_func
    return f(get_current_context(), *args, **kwargs)
  File "/path/.local/lib/python3.9/site-packages/mlem/cli/main.py", line 283, in inner
    res = f(*iargs, **ikwargs) or {}
  File "/path/.local/lib/python3.9/site-packages/mlem/cli/import_object.py", line 46, in import_object
    import_object(
  File "/path/.local/lib/python3.9/site-packages/mlem/api/commands.py", line 399, in import_object
    meta = ImportHook.load_type(type_).process(
  File "/path/.local/lib/python3.9/site-packages/mlem/core/import_objects.py", line 81, in process
    meta = MlemModel.from_obj(data, **kwargs)
  File "/path/.local/lib/python3.9/site-packages/mlem/core/objects.py", line 606, in from_obj
    mt = ModelAnalyzer.analyze(model, sample_data=sample_data)
  File "/path/.local/lib/python3.9/site-packages/mlem/core/model.py", line 336, in analyze
    return super().analyze(obj, sample_data=sample_data, **kwargs)
  File "/path/.local/lib/python3.9/site-packages/mlem/core/hooks.py", line 107, in analyze
    return cls._find_hook(obj).process(obj, **kwargs)
  File "/path/.local/lib/python3.9/site-packages/mlem/core/hooks.py", line 135, in _find_hook
    raise HookNotFound(
mlem.core.errors.HookNotFound: No suitable ModelHook for object of type "OrderedDict". Registered hooks: [<class 'mlem.contrib.callable.CallableModelType'>, <class 'mlem.contrib.torch.TorchModel'>]

[mike0sv]

Hm, it seems that whatever torch loaded was not a torch.nn.Module, but an OrderedDict. How did you save it in the first place?

[schallerala]

Ohh right, you are correct to point that out. Sorry for not raising the matter.

I was trying to reproduce the work of Flowerfan/SF-Net, and the persistence is done like so main.py#L134-L138:

torch.save(model.state_dict(), '%s/%s.%d.pkl' % (args.model_dir, args.model_name, itr))

But I guess it is the "standard"/recommanded way, if we read pytorch's documentation

What do you suggest? Is it part of the limitation of mlem, as out of its use cases?

[mike0sv]

we don't support state_dict yet unfortunately (#332 (comment))
for now you can try saving the whole model

[schallerala]

I guess then that the issue can be closed and merge the PR then. Most likely no need to put more effort into it, with the tests you have implemented.

Thank you for your help and sorry for the struggle today.

Cheers

[mike0sv]

NP!