Learning from pkl files

pkl
pytorch
fastai
Author

Deven Mistry

Published

November 15, 2022

When we train a model using fastai and export it using learn.export, you need to re-declare the functions that you had used to train the model.

Let me show you what I mean by that,

from fastai.vision.all import *
set_seed(99, True)
path = untar_data(URLs.PETS)/'images'

def label_func(x): return x[0].isupper()

dls = ImageDataLoaders.from_name_func(
    path, get_image_files(path), valid_pct=0.2,
    label_func=label_func, item_tfms=Resize(224))

learn = vision_learner(dls, resnet34, metrics=error_rate)
learn.fine_tune(1)
learn.export('saved-model.pkl')

Here, we are training a simple model for image classification and then exporting the saved model.

Now, if you were to use this model in a different jupyter notebook or a py file using the load_learner function.

from fastai.vision.all import *
learn = load_learner('saved-model.pkl')

You would get an error something like this…

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-3-c010bc50794d> in <module>
----> 1 learn = load_learner('saved-model.pkl')

3 frames
/usr/local/lib/python3.7/dist-packages/torch/serialization.py in find_class(self, mod_name, name)
   1040                     pass
   1041             mod_name = load_module_mapping.get(mod_name, mod_name)
-> 1042             return super().find_class(mod_name, name)
   1043
   1044     # Load the data (which may in turn use `persistent_load` to load tensors)

AttributeError: Custom classes or functions exported with your `Learner` not available in namespace.\Re-declare/import before loading:
    Can't get attribute 'label_func' on <module '__main__'>

This error is essentially saying that, before you load your model, the script (py file) or the notebook is looking for the function label_func.

Previously, I would just copy that whole function again from the previous file and paste it in my new file, something like this.

from fastai.vision.all import *
def label_func(x): return x[0].isupper()

learn = load_learner('saved-model.pkl')

But, as it turns out, the load_learner is just looking for a reference to the label_func and not it’s entire definition. So, in theory, you could have your code look something like…

from fastai.vision.all import *

## notice the empty function declaration
def label_func(x): pass

learn = load_learner('saved-model.pkl')

With this declaration, your inference would work just fine and you won’t get any problems.

Note

Use this method only, when you wish to do inference with your model.
If you wish to re-train the model, you would need a new dataloader, label_func and the whole nine-yards.

I hope this was helpful. Happy Learning. Cheers!