Quantization¶

Quantize and export PyTorch models for DeepGate edge deployment.

Functions:

Name	Description
`enable_quantization`	Prepare a trained model for quantization-aware training.
`freeze_quantization`	Freeze a prepared model's measured ranges into a quantized graph.
`post_training_quantize`	Quantize a trained float model by calibrating it on representative data.

dg.enable_quantization ¶

enable_quantization(
    model: Module,
    example_inputs: Tensor | tuple[Tensor, ...],
    quantizer: object | None = None,
    *,
    dynamic_batch: bool = True
) -> torch.fx.GraphModule

Prepare a trained model for quantization-aware training.

Traces the model and inserts simulated quantization, so the returned graph trains with quantization error in the loop. Run your QAT steps (or calibration forwards) on it, then call freeze_quantization.

A logic model has nothing to quantize; use trace for those.

Parameters:

Name	Type	Description	Default
`model`	`Module`	The trained float model to quantize.	required
`example_inputs`	`Tensor \| tuple[Tensor, ...]`	Example inputs to trace with. Pass a batch of at least 2 samples when `dynamic_batch` is set.	required
`quantizer`	`object \| None`	Quantization policy. Defaults to `DGQuantizer`: int8 activations and per-channel int8 weights.	`None`
`dynamic_batch`	`bool`	Let the traced graph accept any batch size.	`True`

Returns:

Type	Description
`GraphModule`	The prepared graph in train mode, on the same device as `model`.

dg.freeze_quantization ¶

freeze_quantization(
    prepared: GraphModule,
) -> torch.fx.GraphModule

Freeze a prepared model's measured ranges into a quantized graph.

Run this after QAT or calibration, once the model has seen enough data to measure its activation ranges.

Parameters:

Name	Type	Description	Default
`prepared`	`GraphModule`	The graph returned by `enable_quantization`.	required

Returns:

Type	Description
`GraphModule`	The quantized graph, ready for `dg.export`.

dg.post_training_quantize ¶

post_training_quantize(
    model: Module,
    dataset: Dataset,
    num_samples: int | None = None,
    *,
    batch_size: int = 128,
    device: str | device | None = None
) -> torch.fx.GraphModule

Quantize a trained float model by calibrating it on representative data.

The whole PTQ path in one call: it measures activation ranges by running samples from dataset through the model, then freezes the result.

model itself comes back unchanged, so the float accuracy you measure after this call is the same one you measured before it. If quantized accuracy falls short, quantization-aware train with enable_quantization and freeze_quantization instead.

Parameters:

Name	Type	Description	Default
`model`	`Module`	The trained float model to quantize.	required
`dataset`	`Dataset`	Representative data to calibrate on. The example input used for tracing is taken from here too.	required
`num_samples`	`int \| None`	Samples to calibrate on. Defaults to the whole dataset.	`None`
`batch_size`	`int`	Samples per calibration forward pass.	`128`
`device`	`str \| device \| None`	Device to calibrate on. Defaults to the model's current device, so a model already on the GPU is calibrated there.	`None`

Returns:

Type	Description
`GraphModule`	The quantized graph, ready for `dg.export`.