Picklability problem when exporting wgan learner

ysig · February 4, 2020, 6:14pm

I tried to save the model trained in:

fastai/course-v3/blob/master/nbs/dl1/lesson7-wgan.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%reload_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from fastai.vision import *\n",

This file has been truncated. show original

by doing learn.export()

Exporting model at: model
Traceback (most recent call last):
  File "fastai_wgan.py", line 40, in <module>
    learn.export(args.model)
  File "/home/syro/miniconda3/envs/fastai/lib/python3.6/site-packages/fastai/basic_train.py", line 243, in export
    try_save(state, self.path, file)
  File "/home/syro/miniconda3/envs/fastai/lib/python3.6/site-packages/fastai/torch_core.py", line 416, in try_save
    torch.save(state, target)
  File "/home/syro/miniconda3/envs/fastai/lib/python3.6/site-packages/torch/serialization.py", line 260, in save
    return _with_file_like(f, "wb", lambda f: _save(obj, f, pickle_module, pickle_protocol))
  File "/home/syro/miniconda3/envs/fastai/lib/python3.6/site-packages/torch/serialization.py", line 185, in _with_file_like
    return body(f)
  File "/home/syro/miniconda3/envs/fastai/lib/python3.6/site-packages/torch/serialization.py", line 260, in <lambda>
    return _with_file_like(f, "wb", lambda f: _save(obj, f, pickle_module, pickle_protocol))
  File "/home/syro/miniconda3/envs/fastai/lib/python3.6/site-packages/torch/serialization.py", line 332, in _save
    pickler.dump(obj)
AttributeError: Can't pickle local object 'AvgFlatten.<locals>.<lambda>'

By trying to make a library that I develop picklable, replacing lambdas with local functions as def(): ... solves this problem.
Is there a way of bypassing this, to save my wgan?

Thanks.