Convenience functions in Python for saving Keras models directly to S3

Keras is a very popular framework developed by Google for training and using machine learning models, and it has become somewhat ubiquitous in its use within the domain. Now I’m no data scientist or machine learning expert, but in my work I am presented with problems related to building things that make machine learning and its related applications easy for data scientists to use.

Serialized machine learning models are almost binary files, making them not the best to store and version control using conventional version control systems such as git. The solution for this is to put them in an object store such as AWS S3 where they can be stored, updated and used by different data scientists on the same team. However, Keras by default stores its models on a folder structure. Take the simple model:

import numpy as np
from tensorflow import keras
inputs = keras.Input(shape=(32,))
outputs = keras.layers.Dense(1)(inputs)
model = keras.Model(inputs, outputs)
model.compile(optimizer="adam", loss="mean_squared_error")
model.save("my_model")

which generates the following folder structure:

Easy folder structure generated using https://tree.nathanfriend.io/

which is only a very simple example of the various folders that Keras could generate depending on the type of model you are creating. By giving the top-level folder name in the fashion of:

model = keras.models.load_model("my_model")

you would be able to load the model for use later.

Problem: Enable easy export of Keras models to S3 without needing to traverse through the generated folder structure in code, and enable easy fetching of a model exported in such a manner so that it can be immediately loaded by Keras.

Solution: Zip up the folder structure generated by Keras in a temporary folder. Upload the zipped file to S3. When loading a model, download the corresponding zip file from S3 in to a temporary folder, unzip it, and load it from there.

Gist for the complete code.

Let’s say we have a simple Keras model like what was outlined above:

inputs = keras.Input(shape=(32,))
outputs = keras.layers.Dense(1)(inputs)
model = keras.Model(inputs, outputs)
model.compile(optimizer="adam", loss="mean_squared_error")

We’re going to use Python’s tempfile library to save this model in a temporary location:

with tempfile.TemporaryDirectory() as tempdir:
    model.save(f"{tempdir}/{model_name}")

By using the temporary directory with context, with tempfile.TemporaryDirectory() , we ensure that the temporary directory is deleted and forgotten as soon as we leave that context block.

Next, we zip it up:

zipf = zipfile.ZipFile(f"{tempdir}/{model_name}.zip", "w", zipfile.ZIP_STORED)
zipdir(f"{tempdir}/{model_name}", zipf)
zipf.close()

This uses a zipdir function which traverses the folder with the Keras model in it, and adds it to the given zip file:

def zipdir(path, ziph):
  # Zipfile hook to zip up model folders
  length = len(path)
  for root, dirs, files in os.walk(path):
    folder = root[length:] # Stop zipping parent folders
    for file in files:
      ziph.write(os.path.join(root, file), os.path.join(folder, file))

Now, we can use an s3fs object to write the zipped file to the S3 bucket we need:

s3fs = s3fs.S3FileSystem(key=AWS_ACCESS_KEY, secret=AWS_SECRET_KEY)
s3fs.put(f"{tempdir}/{model_name}.zip", f"{BUCKET_NAME}/{model_name}.zip")

To get this file back and use it in Keras, we have a simple function that uses all the above libraries to reverse the process:

def s3_get_keras_model(model_name: str) -> keras.Model:
  with tempfile.TemporaryDirectory() as tempdir:
    s3fs = get_s3fs()
    # Fetch and save the zip file to the temporary directory
    s3fs.get(f"{BUCKET_NAME}/{model_name}.zip", f"{tempdir}/{model_name}.zip")
    # Extract the model zip file within the temporary directory
    with zipfile.ZipFile(f"{tempdir}/{model_name}.zip") as zip_ref:
        zip_ref.extractall(f"{tempdir}/{model_name}")
    # Load the keras model from the temporary directory
    return keras.models.load_model(f"{tempdir}/{model_name}")

Put everything together, and we have a simple implementation of saving Keras models in their entirety to S3 and getting them back without having to think about traversing nested folder structures created when saving Keras models.