Saving and loading neural networks is always a little tricky. The best way to do it depends on what exactly you’re trying to do. Do you want to continue training the model? If so, you’ll need to save the optimizer state. If you just want to run it for inference, you might not need this. It also gets more complicated with custom functions. In this post, I’ll walk through how to save a FastAI model and then load it again for inference.

Table of Contents


Let’s train a simple model straight from the FastAI tutorial.

from import *
path = untar_data(URLs.PETS)
files = get_image_files(path/"images")
def label_func(f):
    return f[0].isupper()
dls = ImageDataLoaders.from_name_func(path, files, label_func, item_tfms=Resize(224))
Due to IPython and Windows limitation, python multiprocessing isn't available now.
So `number_workers` is changed to 0 to avoid getting stuck
# if you're on Windows, you need to add this
from PIL import ImageFile
learn = cnn_learner(dls, resnet34, metrics=error_rate)
epoch train_loss valid_loss error_rate time
0 0.818718 1.532427 0.527273 00:08
epoch train_loss valid_loss error_rate time
0 0.300550 0.545057 0.254545 00:05


That looks good. Now let’s save it.


There are two options for saving models in FastAI, and learn.export. saves the model and, by default, also saves the optimizer state. This is what you want to do if you want to resume training.

If you used, you’ll need to use learner.load to load it. You’ll need all of your functions from the previous notebook again.


If you’re going to save it for inference, I recommend using .export(). You should also make sure the model is going to save in the right place by setting learn.path.

learn.path = Path(os.getenv('MODELS'))

If you have custom functions as we do above with label_func, you’ll need to use the dill module to pickle your learner. If you use the default one, it will save the function names but not the content.

import dill
learn.export('catsvdogs.pkl', pickle_module=dill)


Nothing in this section will use any of the code above. This could be done in a completely separate file or notebook.

import dill
from fastai.tabular.all import *

When you load a model, you can load it to the GPU or CPU. By default, it will load to the CPU. If you want it to load to the GPU, you’ll need to pass cpu=False.

learn = load_learner(Path(os.getenv('MODELS')) / 'catsvdogs.pkl', cpu=False, pickle_module=dill)
test = np.zeros((100, 100))
('True', tensor(1), tensor([0.3902, 0.6098]))