.CKPT

Introduction

A checkpoint is a snapshot of a model taken during training, so that training can stop and carry on later, or the best version can be kept. As well as the weights it usually holds the optimiser’s state, the epoch and step counts and the settings used.

The .ckpt extension is used by PyTorch Lightning, which writes one each time it saves a model, and was the usual extension for the original Stable Diffusion image models. Both are files written by PyTorch’s torch.save, in the same layout as a PyTorch model. Like those, they are loaded with Python’s pickle module and can run code, which is why image models moved to safetensors.

Structure

A PyTorch Lightning checkpoint is a dictionary saved with torch.save. Its keys include state_dict (the weights), epoch, global_step, pytorch-lightning_version, optimizer_states and lr_schedulers, and hyper_parameters when the model saved them. A Stable Diffusion .ckpt is a similar dictionary whose state_dict holds the weights of the whole pipeline.

TensorFlow uses the name differently. A TensorFlow checkpoint is a set of files that share a prefix such as model.ckpt: an .index file listing the variables and one or more .data-00000-of-00001 shards holding their values. TensorFlow 1 also wrote a .meta file with the graph.

Links