File size: 5,697 Bytes
477da44 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 |
# 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:
```python
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:
```shell
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:
```shell
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].
```
|