Tensorflow models through a colab submission

I’ve been exclusively training and submitted via google colab notebooks and encounter some road bumps when trying to get tensorflow models submitted. I’m sure this isn’t best practice (please feel free to critique - I’d like to learn better ways!), but thought I’d at least share the outline of how I managed to succeed (there’s been a few related posts already, but I couldn’t completely piece everything together from those).


  • tensorflow 2.4 isn’t supported in the submission environment
  • all “model” components need to be packaged into a single file (zip, pickle, other) - this was slightly painful as I had previously dumped all models/pipelines into dicts and pickled them, but tensorflow models won’t save this way

My solution to tensorflow versioning:

  1. add the following to your ADDITIONAL_PACKAGES config section:

‘tensorflow_addons==0.12.1’, # I wasn’t explicit on the install, but checked what version loaded
‘h5py==2.10.0’ # I wasn’t explicit on the install, but checked what version loaded

  1. pip install the specific version of tensorflow you want (note, you’ll need to restart the runtime afterwards) - I’ve included the bottom four lines below to show you where it goes relative to other code

!pip install tensorflow==2.3
!pip install tensorflow_addons
import importlib
import global_imports
from global_imports import * # do not change this

My solution to model packaging:

  • keep pre-processing in an explicit sklearn pipeline (or dict of multiple pipelines); save tf models to their own separate dict
  • then I zip everything during the save / unzip and put back into dicts during the load

I’ve attached my save an load functions (screenshots only sorry - pasting code in here lost all the nesting)