Skip to main content

tpu_model Model operation tool

The tpu_model tool allows you to view the parameter information of a bmodel file, split a multi-network bmodel into multiple single-network bmodels, and merge multiple network bmodels into a single bmodel.

Usage:
  tpu_model
    --info model_file : show brief model info
    --print model_file : show detailed model info
    --weight model_file : show model weight info
    --update_weight dst_model dst_net dst_offset src_model src_net src_offset
    --extract model_file : extract one multi-net bmodel to multiple one-net bmodels
    --combine file1 .. fileN -o new_file: combine bmodels into one bmodel by file path
    --combine_dir dir1 .. dirN -o new_dir: combine bmodels into one bmodel by directory path
    --dump model_file start_offset byte_size out_file: dump binary data to a file from the bmodel
    --version show tool version
    --kernel_dump model_file -o kernel_file_name : dump kernel module file
    --kernel_update model_file kernel_name : add/update kernel module file
    --custom_ap_update model_file libcpuop_file : add/update custom libcpuop file

Usage​

Currently, the tool supports the following six operations:

View Brief bmodel Information​

tpu_model --info xxx.bmodel

View Detailed bmodel Parameters​

tpu_model --print xxx.bmodel

Split bmodel​

tpu_model --extract xxx.bmodel

This command extracts a multi-network, multi-stage bmodel into separate bmodels, each containing only one network and one stage. The extracted bmodels are named according to their network and stage indices, e.g., bm_net0_stage0.bmodel, bm_net1_stage0.bmodel, etc.

Merge bmodels​

tpu_model --combine a.bmodel b.bmodel c.bmodel -o abc.bmodel

This command merges multiple bmodels into a single bmodel. The -o flag specifies the output file name. If omitted, the default name is compilation.bmodel.

After merging multiple bmodels:

  • If the bmodels have different net_name, the interface will select the corresponding network for inference based on the net_name.
  • If the bmodels have the same net_name, the merged bmodel will support multiple stages (i.e., different input shapes). The interface selects the appropriate stage based on the input shape.
    • For static networks, it chooses the stage that exactly matches the shape.
    • For dynamic networks, it selects the closest matching stage.

Note: When merging bmodels with the same net_name, all must be either statically compiled or dynamically compiled. The tool does not currently support merging static and dynamic compilations of the same net_name.

Merge Folders​

tpu_model --combine_dir a_dir b_dir c_dir -o abc_dir

Similar to the combine function, but it merges entire directories. In addition to merging bmodels, this function also merges input and output test files. The specified folders must contain three essential files generated by the compiler:

  • input_ref_data.dat
  • output_ref_data.dat
  • compilation.bmodel

Export Binary Data​

tpu_model --dump xxx.bmodel start_offset byte_size out_file

This command extracts binary data from a bmodel and saves it to a file. You can use the print command to view the [start, size] values of all binary data, which correspond to the start_offset and byte_size parameters in this command.