ONNX Edit

RKNN-Toolkit2 has a built-in rknn.utils.onnx_edit tool. Used to adjust the onnx model. For some onnx models with special structures, better inference performance can be obtained after the adjusted model is converted to RKNN. The interface is defined as follow:

onnx_edit(model: str,
          export_path : str,
          inputs_transform : dict = {},
          outputs_transform : dict = {}, 
          dataset: Optional[str] = None):
"""
    Args:
        model: The path to the onnx file.
        export_path: The path to the exported onnx file.
        inputs_transform: The dictionary of inputs transform equation.
        outputs_transform: The dictionary of outputs transform equation.
        dataset: The path to the dataset file. dataset file should meet rknn.build requirements. If given, transform of inputs will also apply on dataset and generate a new dataset on the same folder of export_path.
"""

Currently the following functions are available:

Reshape and transpose adjustments to input and output（Modification of dynamic shapes is currently not supported.）

1.Reshape and transpose adjustments to input and output

If some models can reshape and transpose the input and output to a certain extent before converting to RKNN, the graph optimization function of RKNN-Toolkit2 can be used to make the final RKNN model more concise and have better inference performance. For example, by executing the test.py script in this directory, you can compare the differences in the RKNN models generated before and after onnx_edit.

Transformation equation rules:

It must be composed of two parts of symbolic shape, separated by '->'. The left side is the original input symbolic shape, and the right side is the transformed symbolic shape.
The character in symbolic shape are only allowed to be [a-z] or ',' '1'. Other numeric shape characters are not supported except for '1'. The reason is that the same numeric characters cannot determine the context of transpose.
Characters [a-z] are considered independent and have no order relationship, that means, 'a,b->b,a' and 'c,a->a,c' represent the same transformation.
Separated by ',', original input symbolic shape on the left side of the equation must match the dimensions of the input. For example, the original input definition is [32,4,1,64], and the input characters can be 'a,b,c,d' or 'a,b,1,d'
Every character of the original input symbolic shape, except '1', must exist in the transformed symbolic shape. For example, 'a,1,c,d->ac,d,1' is valid, 'a,1, c,d->ac,1' is invalid because 'd' is missing.
The transformed symbolic shape allows the insertion of any number of '1' to achieve a dimension expansion, such as 'a,b,c-> a,1,cb,1,1'
The original input symbolic shape allows using multiple characters and assignment formulas to split the shape. For example, the original input definition is [32,4,1,64], 'ab,c,d,qk[a=2,k =8]->aq,cd,1,kb', which means split 32 into 2x16, split 64 into 8x8, and then perform transpose and reshape operations. The part '[ ]' is called the assignment formula, and multiple formulas are separated by ','. It is allowed that one character in the split is not assigned a value. At this time, the corresponding shape will be automatically inferred. For example, the assignment formula only gives a=2, and ab=32 is known, then b=16 is automatically inferred; if the inferred shape is abnormal An error will be reported directly, for example, ab=32. If a=5 is assigned, then b=6.4. The dimension must be an integer, and an exception error will be thrown.
Support dynamic shape
If the dimensions of a dynamic shape are split, only one of the split dimensions is allowed to be a dynamic shape. Assuming there are 3-dimensional dynamic shapes, ab,c,d->a,c,d,b is invalid , because both a and b are unknown dimensions and cannot be split. And ab,c,d[a=2]->a,c,d,b is valid
Allows assignment of dimensions to dynamic shapes, in which case the corresponding dimensions will become fixed shapes. Assuming there is a 3-dimensional dynamic shape, ab,c,d[a=2]->a,c,d,b will cause the changed dimensions to be [2,c,d,b], and the first dimension will be 2, and ab,c,d[a=2,b=5,c=7,d=11]->a,c,d,b will make the dimension fixed to [2,7,11,5].

Example of transform equation:

1. Reshape 3D to 4D:                'a,b,c->a,b,1,c' or 'a,b,c->a,1,b,c'
2. Reshape 5D to 4D:                'a,b,c,d,e->ab,c,d,e' or 'a,b,c,d,e->a,bce,1,1'
3. Transpose(0,3,1,2):              'a,b,c,d->a,d,b,c'
4. Transpose, merge dim:            'a,b,c,d->ad,c,b' or 'a,b,c,d->d,acb,1'
5. Split dim, transpose, merge dim: 'a,bc,de,f[b=2,d=4]->ab,fe,dc,1',  		#  'c','e' will be auto infered
6. Split dim, transpose, merge dim: 'a,bc,de,f[b=2,c=24,d=4]->ab,fe,dc,1',  #  'bc' = 24 shuold be equal to tensor's 2nd-dim, 'e' will be auto infered

After running test.py, onnx_edit interface will show transform log as follow:

I For 'k_cache.1':'a,b,c,d->1,ad,b,c'
  Input:'k_cache.1' was reset as shape-[1, 32, 48, 1].(Origin shape is [4, 48, 1, 8])
  Insert ops to transform [1, 32, 48, 1] to [4, 48, 1, 8]:
  - Insert reshape op. [1, 32, 48, 1] reshape to [4, 8, 48, 1].
  - Insert transpose op. [4, 8, 48, 1] transpose(0, 2, 3, 1) to [4, 48, 1, 8].
I For 'k_cache':'a,b,c,d->1,ab,c,d'
  Output:'k_cache' was reset as shape-[1, 32, 48, 1].(Origin shape is [4, 8, 48, 1]) 
  Insert ops to transform [4, 8, 48, 1] to [1, 32, 48, 1]:
  - Insert reshape op. [4, 8, 48, 1] reshape to [1, 32, 48, 1].