Failing to load a learner with a custom loss

fastai
1

Hi,

I’m using fastai 1.0.51.
I’ve trained a tabular learner with a custom loss function. I’m passing the loss function like this:

 learn = tabular_learner(
        data, layers=[200, emb_sz], ps=[0.005, 0.05], emb_drop=0.02,
        y_range=y_range, metrics=[rmse],
        callback_fns=ShowGraph, loss_func=custom_loss)

The learner trains correctly and then I’m exporting the model:

 learner.save('tab_learner')
 learner.export('tab_learner')

Then, when I try to load the learner:

learn = ft.load_learner(path=path, file='tab_learner',
                        test=ft.TabularList.from_df(dataset, path=path))

I’m getting the following error:

  File "predict.py", line 43, in main
    preds = get_predictions(test_dataset, model, adjust_by_e=True)
  File "/src/train_embeddings.py", line 129, in get_predictions
    loss_func=custom_loss)
  File "/lib/python3.6/site-packages/fastai/basic_train.py", line 592, in load_learner
    state = torch.load(source, map_location='cpu') if defaults.device == torch.device('cpu') else torch.load(source)
  File "/lib/python3.6/site-packages/torch/serialization.py", line 368, in load
    return _load(f, map_location, pickle_module)
  File "/lib/python3.6/site-packages/torch/serialization.py", line 542, in _load
    result = unpickler.load()
AttributeError: Can't get attribute 'custom_loss' on <module '__main__' from 'predict.py'>

Are there anything wrong in my code? Could it be a bug?

Thank you in advance.

1 Like
2

It’s not a bug, you need to have the code of your custom loss function when you load the Learner run. We can’t serialize the code, just the objects.

1 Like
3

hey @adrianicv, it would be super useful for me if you could paste your loss_function here :slight_smile: