昇腾社区首页
中文
注册

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。

返回值说明

量化后的模型。