torsh-fx 0.1.2

Graph-based model representation and transformation for ToRSh
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28

//! Quantization-aware training (QAT) utilities

use super::context::QuantizationContext;
use super::types::QuantizationScheme;
use crate::{FxGraph, TorshResult};

/// Quantization-aware training (QAT) utilities
pub struct QATUtils;

impl QATUtils {
    /// Prepare a model for QAT
    pub fn prepare_qat(
        graph: &mut FxGraph,
        scheme: QuantizationScheme,
    ) -> TorshResult<QuantizationContext> {
        let mut context = QuantizationContext::new(scheme);
        context.prepare_qat(graph)?;
        Ok(context)
    }

    /// Convert QAT model to quantized model
    pub fn convert_qat_to_quantized(
        graph: &mut FxGraph,
        context: &QuantizationContext,
    ) -> TorshResult<()> {
        context.quantize_graph(graph)
    }
}