quantize¶
- torchao.quantization.quantize_(model: Module, config: AOBaseConfig, filter_fn: Optional[Callable[[Module, str], bool]] = None, device: Optional[Union[device, str, int]] = None)[源代码]¶
使用 config 转换模型中线性模块的权重，模型将被就地修改
- 参数:
model (torch.nn.Module) – 输入模型
config (AOBaseConfig) – 工作流配置对象。
filter_fn (Optional[Callable[[torch.nn.Module, str], bool]]) – 一个函数，它接收一个 nn.Module 实例和该模块的完全限定名称，如果我们要对该模块运行 config，则返回 True
module (权重) –
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))
