bitorch_engine.utils.model_helper.save_checkpoint

bitorch_engine.utils.model_helper.save_checkpoint(model: Module, name: str, qweight_only: bool = True) → None

Saves the state of a quantized PyTorch model in a bit-packed format. This function is intended for models that incorporate quantized layers, allowing for efficient storage and potential speedups in model loading and inference.

The function first packs the layers of the model based on the quantization status of the weights and then saves the model’s state dictionary. The saved checkpoint can be used for inference or to resume training, depending on the inclusion of unpacked weights.

Parameters:

• model (torch.nn.Module) – The model to save. This model should use quantized layers.

• name (str) – The file path where the model checkpoint will be saved. This path should include the filename and the desired file extension (usually “.pth” for PyTorch models).

• qweight_only (bool, optional) – A flag to indicate whether to save only the quantized weights (True) or to also include the original, unpacked weights (False). Saving only quantized weights reduces file size but may limit the ability to resume training. Defaults to True, optimizing for reduced storage.

Returns:

None