Load from SaveModelCallback pickle error

fastai
1

I’m testing out Weights&Biases (wandb.com)'s fastai integrations. I’m following this example code and I’ve added two the two custom callbacks, WandbCallback and LogImagesCallback.

Rather than add them to the callback_fns, in order to exclude these during lr_find, I’ve added them to the fit calls like this:

callbacks = [TerminateOnNaNCallback(), 
                   SaveModelCallback(learn, every='improvement', monitor=monitor, 
                                                    name=best_save_name),
                   WandbCallback(learn), 
                   LogImagesCallback(learn)]
learn.fit_one_cycle(cycles, max_lr, callbacks=callbacks)

# Load best frozen weights, unfreeze, then continue training
learn.load(best_save_name) # fails, pickle error
learn.unfreeze()
learn.fit_one_cycle(cycles, slice(lr), callbacks=callbacks)

Now the learn.load call in between frozen and unfrozen training fails with the error:

Can't pickle local object 'TorchGraph.hook_torch_modules.<locals>.backward_hook'

I’m not yet sure what it is about these custom callbacks that causes this failure. Has anyone experienced a similar issue and found a workaround?

2

Hi, could you try again? The module had been updated since then.
You can add it like any other callback and it shouldn’t run during lr_find.