Can't pickle local object

fastai
1

I am getting an unusual error when I try to export a learner.

Code to make learner and data:

data = (PointsItemList.from_df(df, path='./', cols='image')
        .split_by_valid_func(is_validation)
        .label_from_func(get_coords)
        .transform(tfms, tfm_y=True, size=300, padding_mode='reflection', remove_out=False)
        .databunch(bs=32)
        .normalize(imagenet_stats)
        )

learner = cnn_learner(
    data,
    models.resnet18,
    pretrained=False
    )

After the model has been trained, I run:

learner.save('stage1',return_path=True)
learner.load('stage1')

Which works fine. Then I run:

learner.export('model.pkl')

This is the error I get:

AttributeError Traceback (most recent call last)

[<ipython-input-15-de10b3ec63b2>](https://z5ru7z1tl8o-496ff2e9c6d22116-0-colab.googleusercontent.com/outputframe.html?vrz=colab-20200422-085601-RC00_307797972#) in <module>() ----> 1 learner.export('model.pkl')

3 frames

[/usr/local/lib/python3.6/dist-packages/fastai/basic_train.py](https://z5ru7z1tl8o-496ff2e9c6d22116-0-colab.googleusercontent.com/outputframe.html?vrz=colab-20200422-085601-RC00_307797972#) in export(self, file, destroy) 241 state['data'] = self.data.valid_ds.get_state(**xtra) 242 state['cls'] = self.__class__ --> 243 try_save(state, self.path, file) 244 if destroy: self.destroy() 245

[/usr/local/lib/python3.6/dist-packages/fastai/torch_core.py](https://z5ru7z1tl8o-496ff2e9c6d22116-0-colab.googleusercontent.com/outputframe.html?vrz=colab-20200422-085601-RC00_307797972#) in try_save(state, path, file) 414 #To avoid the warning that come from PyTorch about model not being checked 415 warnings.simplefilter("ignore") --> 416 torch.save(state, target) 417 except OSError as e: 418 raise Exception(f"{e}\n Can't write {path/file}. Pass an absolute writable pathlib obj `fname`.")

[/usr/local/lib/python3.6/dist-packages/torch/serialization.py](https://z5ru7z1tl8o-496ff2e9c6d22116-0-colab.googleusercontent.com/outputframe.html?vrz=colab-20200422-085601-RC00_307797972#) in save(obj, f, pickle_module, pickle_protocol, _use_new_zipfile_serialization) 326 327 with _open_file_like(f, 'wb') as opened_file: --> 328 _legacy_save(obj, opened_file, pickle_module, pickle_protocol) 329 330

[/usr/local/lib/python3.6/dist-packages/torch/serialization.py](https://z5ru7z1tl8o-496ff2e9c6d22116-0-colab.googleusercontent.com/outputframe.html?vrz=colab-20200422-085601-RC00_307797972#) in _legacy_save(obj, f, pickle_module, pickle_protocol) 399 pickler = pickle_module.Pickler(f, protocol=pickle_protocol) 400 pickler.persistent_id = persistent_id --> 401 pickler.dump(obj) 402 403 serialized_storage_keys = sorted(serialized_storages.keys())

AttributeError: Can't pickle local object 'TorchHistory.add_log_hooks_to_pytorch_module.<locals>.<lambda>'
2

We’ll need more information than that to help. How are you building the Learner? How are you building the Data?

Also, for code please wrap it in three ` 's so it’s readable, like I have done so below:

mycode
3

Thanks! I’ve added more detail.

4

I believe it’s complaining about you’re using a lambda function, but I don’t see one in the code you provided. If you’re using a getx or gety lambda in your data loader, then replace it with a function and you’ll be ok.

5

That’s curious. The only thing I have in the labeling function is

def is_validation(imgpath):
    if imgpath.endswith('_val.jpeg'):
      return True
    else:
      return False
6

Are you using a WandB callback?

TorchHistory is referring to that callback apparently. Check source code here.

If you are, just remove the callback before pickling.

4 Likes
7

Ah, that must be it! Thanks so much! I didn’t consider it because it’s not part of the model creation code.

8

Yes, fastai will pickle the Learner objects which includes callbacks, loss function, metrics, etc.

9

@nickvu how did you detach the WandbCallback? I’m facing a very similar error but with TorchGraph rather than TorchHistory. Just like you, I’m not initialising the Learner with WandbCallback but using it only when calling Learner.fit

When I look at learn.cbs I don’t see WandbCallback there as it detaches itself after fitting.

cc @ilovescience

Thanks for your help!

1 Like
10

Full error message:

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-57-c9af5cd1fcdd> in <module>
----> 1 learn.export(fname=f"{fname}.pkl")

~/miniconda3/envs/stable/lib/python3.8/site-packages/fastai/learner.py in export(self, fname, pickle_protocol)
    538         #To avoid the warning that come from PyTorch about model not being checked
    539         warnings.simplefilter("ignore")
--> 540         torch.save(self, self.path/fname, pickle_protocol=pickle_protocol)
    541     self.create_opt()
    542     if state is not None: self.opt.load_state_dict(state)

~/.local/lib/python3.8/site-packages/torch/serialization.py in save(obj, f, pickle_module, pickle_protocol, _use_new_zipfile_serialization)
    362         if _use_new_zipfile_serialization:
    363             with _open_zipfile_writer(opened_file) as opened_zipfile:
--> 364                 _save(obj, opened_zipfile, pickle_module, pickle_protocol)
    365                 return
    366         _legacy_save(obj, opened_file, pickle_module, pickle_protocol)

~/.local/lib/python3.8/site-packages/torch/serialization.py in _save(obj, zip_file, pickle_module, pickle_protocol)
    464     pickler = pickle_module.Pickler(data_buf, protocol=pickle_protocol)
    465     pickler.persistent_id = persistent_id
--> 466     pickler.dump(obj)
    467     data_value = data_buf.getvalue()
    468     zip_file.write_record('data.pkl', data_value, len(data_value))

AttributeError: Can't pickle local object 'TorchGraph.hook_torch_modules.<locals>.backward_hook'
1 Like
11

Did anyone find a solution for this? As @rsomani95 said, I too could not remove the callback as learn.cbs does not show the WandbCallback

12

use dill e.g.

def save_for_meta_learning(args, ckpt_filename='ckpt.pt'):
    if is_lead_worker(args.rank):
        import dill
        args.logger.save_current_plots_and_stats()
        # - ckpt
        assert uutils.xor(args.training_mode == 'epochs', args.training_mode == 'iterations')
        args_pickable = uutils.make_args_pickable(args)
        # args.meta_learner.args = args_pickable
        f: nn.Module = get_model_from_ddp(args.base_model)
        # pickle vs torch_uu.save https://discuss.pytorch.org/t/advantages-disadvantages-of-using-pickle-module-to-save-models-vs-torch-save/79016
        torch.save({'training_mode': args.training_mode,  # its or epochs
                    'it': args.it,
                    'epoch_num': args.epoch_num,
                    # 'args': args_pickable,
                    'args_pickable': args_pickable,
                    # 'meta_learner': args.meta_learner,
                    'meta_learner_str': str(args.meta_learner),
                    # 'f': f,
                    'f_state_dict': f.state_dict(),
                    'f_str': str(f),
                    # 'f_modules': f._modules,
                    # 'f_modules_str': str(f._modules),
                    'outer_opt_state_dict': args.outer_opt.state_dict()
                    },
                   pickle_module=dill,
                   f=args.log_root / ckpt_filename)