machine learning
model saving
scikit-learn
keras
tensorflow
mxnet

What are all the formats to save machine learning model in scikit-learn, keras, tensorflow and mxnet?

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

Introduction

There is no single model file format that covers every machine learning framework equally well. Each library has its own native persistence story, and some also support interchange formats that are useful only in specific deployment pipelines.

Scikit-learn Formats

Scikit-learn models are Python objects, so the most common persistence methods serialize those objects rather than exporting a framework-specific graph.

Common options are:

  • 'pickle'
  • 'joblib'
  • 'cloudpickle'
  • 'skops.io'
  • ONNX, when a converter exists for the estimator

joblib is common for local Python workflows:

python
1from joblib import dump, load
2from sklearn.ensemble import RandomForestClassifier
3
4model = RandomForestClassifier(random_state=0)
5model.fit([[0, 0], [1, 1]], [0, 1])
6
7dump(model, "model.joblib")
8loaded = load("model.joblib")
9print(loaded.predict([[1, 1]]))

pickle and joblib are simple, but they are not safe for untrusted files. skops.io exists to offer a more controlled persistence story for scikit-learn objects.

Keras Formats

In modern Keras, the recommended whole-model format is the .keras format. It stores architecture, weights, and training configuration together.

python
1import keras
2
3model = keras.Sequential([
4    keras.layers.Input(shape=(4,)),
5    keras.layers.Dense(8, activation="relu"),
6    keras.layers.Dense(1)
7])
8
9model.save("model.keras")
10restored = keras.models.load_model("model.keras")

Weights-only saving is also common:

python
model.save_weights("model.weights.h5")

Legacy HDF5 whole-model files such as model.h5 still appear in older codebases, but .keras is the current default direction for Keras whole-model persistence.

TensorFlow Formats

TensorFlow has two major persistence concepts that people often mix together:

  • checkpoints for training state
  • exported models for inference and serving

Training checkpoints are useful when you want to resume training:

python
1import tensorflow as tf
2
3model = tf.keras.Sequential([
4    tf.keras.layers.Input(shape=(4,)),
5    tf.keras.layers.Dense(2)
6])
7
8ckpt = tf.train.Checkpoint(model=model)
9ckpt.write("checkpoints/ckpt")

For serving or deployment, TensorFlow uses SavedModel-style export artifacts:

python
1import tensorflow as tf
2
3model = tf.keras.Sequential([
4    tf.keras.layers.Input(shape=(4,)),
5    tf.keras.layers.Dense(2)
6])
7
8model.export("exported_model")

In TensorFlow-centric code, you may also encounter Keras-managed .keras files, because Keras models can run on top of TensorFlow. The important distinction is whether you are saving a training object for reuse in Python or exporting an inference artifact for serving.

MXNet Formats

MXNet has a few common persistence patterns depending on the API style you use.

For Gluon parameter-only saving:

python
1from mxnet.gluon import nn
2
3net = nn.Dense(2)
4net.initialize()
5net.save_parameters("model.params")

For HybridBlock export, MXNet commonly writes a symbol file plus parameters:

  • a JSON graph file, often ending in -symbol.json
  • a parameters file, often ending in .params

Older module and checkpoint APIs also save graph and parameter artifacts separately. In other words, MXNet persistence is often multi-file rather than one universal archive.

How to Choose

Choose the format based on how the model will be used after saving:

  • Use joblib or similar Python object persistence for scikit-learn workflows that stay in Python.
  • Use .keras for modern Keras whole-model saves.
  • Use TensorFlow checkpoints to resume training.
  • Use TensorFlow exported models for serving.
  • Use MXNet parameter or checkpoint formats that match the API family you are using.

If interoperability matters, check whether ONNX export is supported for your specific model and operators. That is a deployment decision, not a universal default.

Common Pitfalls

  • Asking for "all formats" without separating training checkpoints, whole-model saves, and inference exports. They solve different problems.
  • Loading untrusted pickle or joblib files. These formats can execute arbitrary code during deserialization.
  • Assuming old HDF5 examples are the current Keras recommendation.
  • Mixing TensorFlow checkpoints and Keras weight files as if they were interchangeable.

Summary

  • Scikit-learn commonly uses pickle, joblib, cloudpickle, skops.io, and sometimes ONNX.
  • Modern Keras recommends the .keras whole-model format, with .weights.h5 for weights only.
  • TensorFlow distinguishes between checkpoints for training and exported models for inference.
  • MXNet commonly uses parameter files, checkpoint files, and graph-plus-params exports.
  • The right format depends on whether you need Python reuse, training resume, serving, or cross-framework interchange.

Course illustration
Course illustration

All Rights Reserved.