Load_learner not working with flask

Hello everyone,

I’m trying to load my fastai model into a Flask web server and I cannot figure out how to solve the following error.
I exported my model with learn.export() and the following script works without error:

# this script works without error
from fastai.vision.all import *

model_path = 'my_model.pkl'
test_img_path = 'test_img.jpg'

def get_focus_point(path_name):
    dfb = next(iter(df[df['name']==path_name.name].index), ('no match for '+path_name.name))
    return tensor([df['x_p'][dfb], df['y_p'][dfb]])

model = load_learner(model_path)

print(model.predict(test_img_path))

get_focus_point is my custom function for get_y of the DataBlock, which seems to be needed to load the model.
When I try to use the exact same code (that worked in the normal python script above) with Flask, so this code:

# this script does not work
from flask import Flask
from fastai.vision.all import *

app = Flask(__name__)

model_path = 'my_model.pkl'

def get_focus_point(path_name):
    dfb = next(iter(df[df['name']==path_name.name].index), ('no match for '+path_name.name))
    return tensor([df['x_p'][dfb], df['y_p'][dfb]])

model = load_learner(model_path)

@app.route('/')
def hello_world():
    return 'Hello, World!'

the following error occurs:

 * Serving Flask app "test_flask.py"
 * Environment: production
   WARNING: This is a development server. Do not use it in a production deployment.
   Use a production WSGI server instead.
 * Debug mode: off
Traceback (most recent call last):
  File "/home/kilian/.local/bin/flask", line 8, in <module>
    sys.exit(main())
  File "/home/kilian/.local/lib/python3.6/site-packages/flask/cli.py", line 967, in main
    cli.main(args=sys.argv[1:], prog_name="python -m flask" if as_module else None)
  File "/home/kilian/.local/lib/python3.6/site-packages/flask/cli.py", line 586, in main
    return super(FlaskGroup, self).main(*args, **kwargs)
  File "/home/kilian/.local/lib/python3.6/site-packages/click/core.py", line 782, in main
    rv = self.invoke(ctx)
  File "/home/kilian/.local/lib/python3.6/site-packages/click/core.py", line 1259, in invoke
    return _process_result(sub_ctx.command.invoke(sub_ctx))
  File "/home/kilian/.local/lib/python3.6/site-packages/click/core.py", line 1066, in invoke
    return ctx.invoke(self.callback, **ctx.params)
  File "/home/kilian/.local/lib/python3.6/site-packages/click/core.py", line 610, in invoke
    return callback(*args, **kwargs)
  File "/home/kilian/.local/lib/python3.6/site-packages/click/decorators.py", line 73, in new_func
    return ctx.invoke(f, obj, *args, **kwargs)
  File "/home/kilian/.local/lib/python3.6/site-packages/click/core.py", line 610, in invoke
    return callback(*args, **kwargs)
  File "/home/kilian/.local/lib/python3.6/site-packages/flask/cli.py", line 848, in run_command
    app = DispatchingApp(info.load_app, use_eager_loading=eager_loading)
  File "/home/kilian/.local/lib/python3.6/site-packages/flask/cli.py", line 305, in __init__
    self._load_unlocked()
  File "/home/kilian/.local/lib/python3.6/site-packages/flask/cli.py", line 330, in _load_unlocked
    self._app = rv = self.loader()
  File "/home/kilian/.local/lib/python3.6/site-packages/flask/cli.py", line 388, in load_app
    app = locate_app(self, import_name, name)
  File "/home/kilian/.local/lib/python3.6/site-packages/flask/cli.py", line 240, in locate_app
    __import__(module_name)
  File "/home/kilian/projects/focusfinder/webserver/test_flask.py", line 24, in <module>
    model = load_learner((models_path+model_name))
  File "/home/kilian/projects/python_packages/fastai/fastai/learner.py", line 549, in load_learner
    res = torch.load(fname, map_location='cpu' if cpu else None)
  File "/home/kilian/.local/lib/python3.6/site-packages/torch/serialization.py", line 594, in load
    return _load(opened_zipfile, map_location, pickle_module, **pickle_load_args)
  File "/home/kilian/.local/lib/python3.6/site-packages/torch/serialization.py", line 853, in _load
    result = unpickler.load()
AttributeError: Can't get attribute 'get_focus_point' on <module '__main__' from '/home/kilian/.local/bin/flask'>

So it basically cannot find my get_focus_point function anymore, that I define at the top of the file.
Now my question, how can I define my function that it can still be found by load_learner() in a Flask web server?

I’d really appreciate any help!

1 Like

Running into similar problems. Have you found any solution, yet?

Unfortunately not, I am pretty sure by now that somehow the function is missing in the unpickle environment but I don´t know yet how to load it to that environment. So it seems like the place of unpickling changes when flask is used.