クラス torch_quantizer は、クオンタイザー オブジェクトを作成します。
class torch_quantizer():
def __init__(self,
quant_mode: str, # ['calib', 'test']
module: torch.nn.Module,
input_args: Union[torch.Tensor, Sequence[Any]] = None,
state_dict_file: Optional[str] = None,
output_dir: str = "quantize_result",
bitwidth: int = 8,
device: torch.device = torch.device("cuda"),
quant_config_file: Optional[str] = None):
):
引数
- quant_mode
- プロセスが使用している量子化モードを指定する整数。「calib」は量子化のキャリブレーション、「test」は量子化されたモデルの評価。
- module
- 量子化される浮動小数点モジュール。
- input_args
- 量子化される浮動小数点モジュールの実際の入力と同じ形状を持つ入力テンソル。値は乱数を取ることができます。
- state_dict_file
- 浮動小数点モジュールのトレーニング済みパラメーター ファイル。浮動小数点モジュールにパラメーターが読み込まれている場合は、このパラメーターを設定する必要はありません。
- output_dir
- 量子化の結果と中間ファイルを格納するディレクトリ。デフォルトは「quantize_result」です。
- bitwidth
- グローバルな量子化ビット幅。デフォルトは 8 です。
- デバイス
- GPU または CPU 上の実行モデル。
- quant_config_file
- 量子化ストラテジの設定を格納する TJson ファイルのパス。