def quantize
函数功能
传入浮点模型,通过导出的量化权重和对应描述符转化为处理量化权重并返回量化模型。
函数原型
def quantize(model, quant_des_path, timestep_config=None, dtype=torch.bfloat16):
参数说明
参数名 |
输入 |
类型 |
说明 |
---|---|---|---|
model |
输入 |
float |
需要为nn.Module。 |
quant_des_path |
输入 |
str |
通过工具导出的权重描述符全路径。 |
timestep_config |
可选输入 |
- |
当使用时间步量化算法时需要输入,需要为class TimestepPolicyConfig类型。 |
dtype |
可选输入 |
torch.float16/torch.bfloat16 |
可选输入,默认torch.bfloat16。 |
返回值说明
量化后的模型。
父主题: API参考(Python)