地瓜官网
地瓜开发者社区
地瓜官网
地瓜开发者社区

quantization.hbdk4.export

horizon_plugin_pytorch.quantization.hbdk4.export (model: Module, example_inputs: Any, *, name: str = 'forward', input_names: Any | None = None, output_names: Any | None = None, input_descs: Any | None = None, output_descs: Any | None = None, native_pytree: bool = True)

Export nn.Module to hbir model.

Parameters:

model (Module) – Input model.

example_inputs (Any) – Example input for tracing.

name (str) – The name of func in exported module. Users can get the func by getattr(hbir_module, name).

input_names (Optional[Any]) – Set hbir inputs with given names, should have the same structure with example_inputs.

output_names (Optional[Any]) – Set hbir outputs with given names, should have the same structure with model output.

input_descs (Optional[Any]) – Set hbir inputs with given descriptions, should have the same structure with example_inputs.

output_descs (Optional[Any]) – Set hbir outputs with given descriptions, should have the same structure with model output.

native_pytree (bool) –

Whether use native pytree support provided by hbdk4. Assume the model input is passed to torch model as:

torch_rets = torch_model(*model_args)

When using hbdk4 pytree, the hbir model can be called as:

hbir_rets = hbir_model.functions[0](model_args) assert torch_rets == hbir_rets

When using plugin pytree, the hbir model can be called as:

flat_inputs = get_hbir_input_flattener(hbir_model)(model_args) flat_rets = hbir_model.functions[0](*flat_inputs) hbir_rets = get_hbir_output_unflattener(hbir_model)(flat_rets) assert torch_rets == hbir_rets

Return type: Module

Returns: Hbir model wrapped with Module.