Failed attempt to serve PyTorch model

[ankxyz]

Train & saving model code:

import torch
import math
from mlem.api import save


x = torch.linspace(-math.pi, math.pi, 2000)
y = torch.sin(x)

p = torch.tensor([1, 2, 3])
xx = x.unsqueeze(-1).pow(p)

model = torch.nn.Sequential(
    torch.nn.Linear(3, 1),
    torch.nn.Flatten(0, 1)
)
loss_fn = torch.nn.MSELoss(reduction='sum')
learning_rate = 1e-6

for t in range(2000):

    y_pred = model(xx)
    loss = loss_fn(y_pred, y)

    if t % 100 == 99:
        print(t, loss.item())

    model.zero_grad()
    loss.backward()

    with torch.no_grad():
        for param in model.parameters():
            param -= learning_rate * param.grad

linear_layer = model[0]

print(f'Result: y = {linear_layer.bias.item()} + {linear_layer.weight[:, 0].item()} x + {linear_layer.weight[:, 1].item()} x^2 + {linear_layer.weight[:, 2].item()} x^3')


save(
    model,
    'torch_model',
    sample_data=xx,
    description='Example of pytorch model'
)

Environment:

• OS: Ubuntu 20.04
• Python: 3.8.10
• Python package manager: pip 22.1.2
• mlem: 0.2.3
• pytorch: 1.11.0
• fastapi: 0.78.0
• uvicorn: 0.17.6

Metafile of mlem model:

artifacts:
  data:
    hash: fa0c8b5bbc44d9a19b12d47850609ac5
    size: 1767
    uri: torch_model
description: Example of pytorch model
model_type:
  methods:
    predict:
      args:
      - name: data
        type_:
          dtype: float32
          shape:
          - null
          - 3
          type: torch
      name: predict
      returns:
        dtype: float32
        shape:
        - null
        type: torch
    torch_predict:
      args: []
      name: _call_impl
      returns:
        dtype: float32
        shape:
        - null
        type: torch
      varargs: input
      varkw: kwargs
  type: torch
object_type: model
requirements:
- module: torch
  version: 1.11.0+cu102

Command:

mlem serve torch_model fastapi

Ouptut/Error:

⏳️ Loading model from .mlem/model/torch_model.mlem
Starting fastapi server...
❌ Unexpected error: 
Please report it here: <https://github.com/iterative/mlem/issues>

Could you please give some simple example how to save and serve PyTorch models by mlem?

[mike0sv]

Please, add mlem --tb serve torch_model fastapi output. We'll take a look

[ankxyz]

Please, add mlem --tb serve torch_model fastapi output. We'll take a look

Command:

mlem --tb serve torch_model fastapi

Output:

⏳️ Loading model from .mlem/model/torch_model.mlem
Starting fastapi server...
Traceback (most recent call last):
  File "/media/alex/hdd/Dev/progexp/mlem/pytorch-example/.venv/bin/mlem", line 8, in <module>
    sys.exit(app())
  File "/media/alex/hdd/Dev/progexp/mlem/pytorch-example/.venv/lib/python3.8/site-packages/typer/main.py", line 214, in __call__
    return get_command(self)(*args, **kwargs)
  File "/media/alex/hdd/Dev/progexp/mlem/pytorch-example/.venv/lib/python3.8/site-packages/click/core.py", line 1128, in __call__
    return self.main(*args, **kwargs)
  File "/media/alex/hdd/Dev/progexp/mlem/pytorch-example/.venv/lib/python3.8/site-packages/click/core.py", line 1053, in main
    rv = self.invoke(ctx)
  File "/media/alex/hdd/Dev/progexp/mlem/pytorch-example/.venv/lib/python3.8/site-packages/click/core.py", line 1659, in invoke
    return _process_result(sub_ctx.command.invoke(sub_ctx))
  File "/media/alex/hdd/Dev/progexp/mlem/pytorch-example/.venv/lib/python3.8/site-packages/click/core.py", line 1395, in invoke
    return ctx.invoke(self.callback, **ctx.params)
  File "/media/alex/hdd/Dev/progexp/mlem/pytorch-example/.venv/lib/python3.8/site-packages/click/core.py", line 754, in invoke
    return __callback(*args, **kwargs)
  File "/media/alex/hdd/Dev/progexp/mlem/pytorch-example/.venv/lib/python3.8/site-packages/typer/main.py", line 500, in wrapper
    return callback(**use_params)  # type: ignore
  File "/media/alex/hdd/Dev/progexp/mlem/pytorch-example/.venv/lib/python3.8/site-packages/click/decorators.py", line 26, in new_func
    return f(get_current_context(), *args, **kwargs)
  File "/media/alex/hdd/Dev/progexp/mlem/pytorch-example/.venv/lib/python3.8/site-packages/mlem/cli/main.py", line 278, in inner
    res = f(*iargs, **ikwargs) or {}
  File "/media/alex/hdd/Dev/progexp/mlem/pytorch-example/.venv/lib/python3.8/site-packages/mlem/cli/serve.py", line 39, in serve
    serve(
  File "/media/alex/hdd/Dev/progexp/mlem/pytorch-example/.venv/lib/python3.8/site-packages/mlem/api/commands.py", line 348, in serve
    server_obj.serve(interface)
  File "/media/alex/hdd/Dev/progexp/mlem/pytorch-example/.venv/lib/python3.8/site-packages/mlem/contrib/fastapi.py", line 100, in serve
    app = self.app_init(interface)
  File "/media/alex/hdd/Dev/progexp/mlem/pytorch-example/.venv/lib/python3.8/site-packages/mlem/contrib/fastapi.py", line 86, in app_init
    handler, response_model = self._create_handler(
  File "/media/alex/hdd/Dev/progexp/mlem/pytorch-example/.venv/lib/python3.8/site-packages/mlem/contrib/fastapi.py", line 50, in _create_handler
    kwargs = {
  File "/media/alex/hdd/Dev/progexp/mlem/pytorch-example/.venv/lib/python3.8/site-packages/mlem/contrib/fastapi.py", line 51, in <dictcomp>
    key: (serializer.get_model(), ...)
  File "/media/alex/hdd/Dev/progexp/mlem/pytorch-example/.venv/lib/python3.8/site-packages/mlem/contrib/torch.py", line 72, in get_model
    raise NotImplementedError
NotImplementedError

[mike0sv]

Ok, that was a missing implementation on our end. Please try installing from #296