The ability to save and restore the state of a model is vital for a number of applications, such as in transfer learning or for performing inference using pretrained models. Saving the parameters of a model (weights, biases, etc.) in a checkpoint file or directory is one way to accomplish this.
This module provides a high-level interface for loading and saving TensorFlow v2 format checkpoints, as well as lower-level components that write to and read from this file format.
Loading and saving simple models
By conforming to the Checkpointable protocol, many simple models can be serialized to
checkpoints without any additional code:
import Checkpoints
import ImageClassificationModels
extension LeNet: Checkpointable {}
var model = LeNet()
...
try model.writeCheckpoint(to: directory, name: "LeNet")
and then that same checkpoint can be read by using:
try model.readCheckpoint(from: directory, name: "LeNet")
This default implementation for model loading and saving will use a path-based naming
scheme for each tensor in the model that is based on the names of the properties within the
model structs. For example, the weights and biases within the first convolution in
the LeNet-5 model
will be saved with the names conv1/filter and conv1/bias, respectively. When loading,
the checkpoint reader will search for tensors with these names.
Customizing model loading and saving
If you want to have greater control over which tensors are saved and loaded, or the naming
of those tensors, the Checkpointable protocol offers a few points of customization.
To ignore properties on certain types, you can provide an implementation of
ignoredTensorPaths on your model that returns a Set of strings in the form of
Type.property. For example, to ignore the scale property on every Attention layer, you
could return ["Attention.scale"].
By default, a forward slash is used to separate each deeper level in a model. This can be
customized by implementing checkpointSeparator on your model and providing a new
string to use for this separator.
Finally, for the greatest degree of customization in tensor naming, you can implement
tensorNameMap and provide a function that maps from the default string name generated
for a tensor in the model to a desired string name in the checkpoint. Most commonly, this
will be used to interoperate with checkpoints generated with other frameworks, each of which
have their own naming conventions and model structures. A custom mapping function gives
the greatest degree of customization for how these tensors are named.
Some standard helper functions are provided, like the default
CheckpointWriter.identityMap (which simply uses the automatically generated tensor
path name for checkpoints), or the CheckpointWriter.lookupMap(table:) function,
which can build a mapping from a dictionary.
For an example of how custom mapping can be accomplished, please see the GPT-2 model, which uses a mapping function to match the exact naming scheme used for OpenAI's checkpoints.
The CheckpointReader and CheckpointWriter components
For checkpoint writing, the extension provided by the Checkpointable protocol
uses reflection and keypaths to iterate over a model's properties and generate a dictionary
that maps string tensor paths to Tensor values. This dictionary is provided to an underlying
CheckpointWriter, along with a directory in which to write the checkpoint. That
CheckpointWriter handles the task of generating the on-disk checkpoint from that
dictionary.
The reverse of this process is reading, where a CheckpointReader is given the location of
an on-disk checkpoint directory. It then reads from that checkpoint and forms a dictionary that
maps the names of tensors within the checkpoint with their saved values. This dictionary is
used to replace the current tensors in a model with the ones in this dictionary.
For both loading and saving, the Checkpointable protocol maps the string paths to tensors
to corresponding on-disk tensor names using the above-described mapping function.
If the Checkpointable protocol lacks needed functionality, or more control is desired over
the checkpoint loading and saving process, the CheckpointReader and
CheckpointWriter classes can be used by themselves.
The TensorFlow v2 checkpoint format
The TensorFlow v2 checkpoint format, as briefly described in this header, is the second generation format for TensorFlow model checkpoints. This second-generation format has been in use since late 2016, and has a number of improvements over the v1 checkpoint format. TensorFlow SavedModels use v2 checkpoints within them to save model parameters.
A TensorFlow v2 checkpoint consists of a directory with a structure like the following:
checkpoint/modelname.index
checkpoint/modelname.data-00000-of-00002
checkpoint/modelname.data-00001-of-00002
where the first file stores the metadata for the checkpoint and the remaining files are binary shards holding the serialized parameters for the model.
The index metadata file contains the types, sizes, locations, and string names of all serialized
tensors contained in the shards. That index file is the most structurally complex part of the
checkpoint, and is based on tensorflow::table, which is itself based on SSTable / LevelDB.
This index file is composed of a series of key-value pairs, where the keys are strings and the
values are protocol buffers. The strings are sorted and prefix-compressed. For example: if
the first entry is conv1/weight and next conv1/bias, the second entry only uses the
bias part.
This overall index file is sometimes compressed using
Snappy compression. The
SnappyDecompression.swift file provides a native Swift implementation of Snappy
decompression from a compressed Data instance.
The index header metadata and tensor metadata are encoded as protocol buffers and encoded / decoded directly via Swift Protobuf.
The CheckpointIndexReader and CheckpointIndexWriter classes handle loading
and saving these index files as part of the overarching CheckpointReader and
CheckpointWriter classes. The latter use the index files as basis for determining what to
read from and write to the structurally simpler binary shards that contain the tensor data.