Fastai v2 Export and Prediction issue

hi All,
I am new to Fastai and am using the binary segmentation notebook that is in the Zach Mueller’s lessons (thanks Zach for amazing work!!).
My model is as follows:
learn = unet_learner(dls, resnet34)
learn.model_dir=‘gdrive/My Drive/ML/Project1/models/’
learn.fit(10, max_lr=slice(1e-5, 1e-3))
learn.save(‘stage1’)

This saves the model in the models directory: gdrive/My Drive/ML/Project1/models/stage1.pth

Now, like FastAi v1, I wanted to export model to create a pkl file that I can load to create a new learner for Prediction.
learn.export()

Here I get an error:

PicklingError Traceback (most recent call last)
in ()
----> 1 learn.export()

4 frames
/usr/local/lib/python3.6/dist-packages/torch/serialization.py in _save(obj, f, pickle_module, pickle_protocol)
294 pickler = pickle_module.Pickler(f, protocol=pickle_protocol)
295 pickler.persistent_id = persistent_id
–> 296 pickler.dump(obj)
297
298 serialized_storage_keys = sorted(serialized_storages.keys())

PicklingError: Can’t pickle <function at 0x7fb30fd85598>: attribute lookup on main failed

My requirement is to crete a new notebook, when I load the exported model and create an inference server that will help me to load 1 single image that will be segmented.

In the class notes, the inference is created in the same notebook as training. Is there any other notebook where I can load the model and run inference?

How are you creating the DataBlock? Is there any lambda functions being used? If so replace them with full functions (def myFunc()) as lambda wont pickle

Thanks Zachary. I defined DataBlock same as you defined in your notebook:
get_y = lambda o: get_msk(o, p2c)

binary = DataBlock(blocks=(ImageBlock, MaskBlock(codes)), get_items=get_image_files, splitter=RandomSplitter(), get_y=get_y, item_tfms=Resize(256),batch_tfms=[Normalize.from_stats(*imagenet_stats)])

dls = binary.dataloaders(path_img, bs=8)
learn = unet_learner(dls, resnet34)
learn.fit(10, max_lr=slice(1e-5, 1e-3))

learn.save(file=‘seg-1705-01.pth’)
learn.export()

Is the problem because of get_y has lambda function? How do I define it?
Sorry I am new to Pytorch and FastAi2, so any help would be great!!

Yes, it is. You should change your lambda function to a regular python function instead, doing the same thing (but without lambda)

2 Likes

Thanks. It worked. I changed the lambda function to a normal python function, as below:
def get_y(o): return get_msk(o,p2c)

binary = DataBlock(blocks=(ImageBlock, MaskBlock(codes)), get_items=get_image_files, splitter=RandomSplitter(), get_y=get_y, item_tfms=Resize(256),batch_tfms=[Normalize.from_stats(*imagenet_stats)])

Rest of the code was as previous.
Thanks a lot for quick response. You are amazing!! :+1:

2 Likes

Hello and thank you for your post.
I’m new to Fastai too, and I want to do like you :

My requirement is to create a new notebook, when I load the exported model and create an inference server that will help me to load 1 single image that will be segmented.

In the class notes, the inference is created in the same notebook as training. Is there any other notebook where I can load the model and run inference?

Now is OK for the export, but when I open a new notebook, and I load my model, I have an error :

Do you have the same problem ?