quantize¶
- torchao.quantization.quantize_(model: Module, config: Union[AOBaseConfig, Callable[[Module], Module]], filter_fn: Optional[Callable[[Module, str], bool]] = None, set_inductor_config: Optional[bool] = None, device: Optional[Union[device, str, int]] = None)[源]¶
使用 config 轉換模型中線性模組的權重，模型會被原地修改
- 引數:
model (torch.nn.Module) – 輸入模型
config (Union[AOBaseConfig, Callable[[torch.nn.Module], torch.nn.Module]]) – (1) 工作流配置物件 或 (2) 一個將 tensor subclass 轉換應用於模組權重並返回該模組的函式（例如，將 linear 模組的權重張量轉換為仿射量化張量）。注意：(2) 將在未來版本中刪除。
filter_fn (Optional[Callable[[torch.nn.Module, str], bool]]) – 接受 nn.Module 例項和模組完全限定名的函式，如果希望對該模組的權重執行 config，則返回 True
module (該模組的權重) –
set_inductor_config (bool, optional) – 是否自動使用推薦的 inductor 配置設定（預設為 None）
device (device, optional) – 在應用 filter_fn 之前將模組移動到的裝置。可以設定為 “cuda” 以加速量化。最終模型將位於指定的 device 上。預設為 None（不改變裝置）。
示例
import torch import torch.nn as nn from torchao import quantize_ # quantize with some predefined `config` method that corresponds to # optimized execution paths or kernels (e.g. int4 tinygemm kernel) # also customizable with arguments # currently options are # int8_dynamic_activation_int4_weight (for executorch) # int8_dynamic_activation_int8_weight (optimized with int8 mm op and torch.compile) # int4_weight_only (optimized with int4 tinygemm kernel and torch.compile) # int8_weight_only (optimized with int8 mm op and torch.compile from torchao.quantization.quant_api import int4_weight_only m = nn.Sequential(nn.Linear(32, 1024), nn.Linear(1024, 32)) quantize_(m, int4_weight_only(group_size=32))
