[TFLite] Support PRelu

[FrozenGene]

As discuss forum topic: https://discuss.tvm.ai/t/missing-tflite-operators/3150/8, we want to have PRelu support in TFLite

@apivovarov

[tqchen]

Thanks @FrozenGene

[apivovarov]

@FrozenGene
I tried to compile my model and got the following unable to unify errors. ( I also added model visualization screenshots at the end)

 %0 = nn.pad(%input_1, pad_width=[[0, 0], [0, 1], [0, 1], [0, 0]]);
%1 = nn.conv2d(%0, %v_param_1, strides=[2, 2], channels=16, kernel_size=[3, 3], data_layout="NHWC", kernel_layout="HWIO");
%2 = nn.bias_add(%1, %v_param_2, axis=3);
%3 = nn.prelu(%2, %v_param_3, axis=3) tensor type `Tensor[(16), float32]` has 1 dimensions, while `Tensor[(1, 1, 16), float32]` has 3 dimensions; unable to unify: `Tensor[(16), float32]` and `Tensor[(1, 1, 16), float32]`; ;
%4 = nn.conv2d(%3, %v_param_4, channels=8, kernel_size=[1, 1], data_layout="NHWC", kernel_layout="HWIO");
%5 = nn.bias_add(%4, %v_param_5, axis=3);
%6 = nn.prelu(%5, %v_param_6, axis=3) tensor type `Tensor[(8), float32]` has 1 dimensions, while `Tensor[(1, 1, 8), float32]` has 3 dimensions; unable to unify: `Tensor[(8), float32]` and `Tensor[(1, 1, 8), float32]`; ;
%7 = nn.pad(%6, pad_width=[[0, 0], [1, 1], [1, 1], [0, 0]]);

 %112 = nn.conv2d(%111, %v_param_89, channels=32, kernel_size=[1, 1], data_layout="NHWC", kernel_layout="HWIO");
%113 = nn.bias_add(%112, %v_param_90, axis=3);
%114 = add(%105, %113);
%115 = nn.prelu(%114, %v_param_91, axis=3) tensor type `Tensor[(32), float32]` has 1 dimensions, while `Tensor[(1, 1, 32), float32]` has 3 dimensions; unable to unify: `Tensor[(32), float32]` and `Tensor[(1, 1, 32), float32]`; ;
%116 = nn.conv2d(%115, %v_param_92, channels=16, kernel_size=[1, 1], data_layout="NHWC", kernel_layout="HWIO");
%117 = nn.bias_add(%116, %v_param_93, axis=3);
%118 = nn.prelu(%117, %v_param_94, axis=3) tensor type `Tensor[(16), float32]` has 1 dimensions, while `Tensor[(1, 1, 16), float32]` has 3 dimensions; unable to unify: `Tensor[(16), float32]` and `Tensor[(1, 1, 16), float32]`; ;
%119 = nn.pad(%118, pad_width=[[0, 0], [1, 1], [1, 1], [0, 0]]);

 %538 = nn.conv2d(%537, %v_param_425, channels=256, kernel_size=[1, 1], data_layout="NHWC", kernel_layout="HWIO");
%539 = nn.bias_add(%538, %v_param_426, axis=3);
%540 = add(%531, %539);
%541 = nn.prelu(%540, %v_param_427, axis=3) tensor type `Tensor[(256), float32]` has 1 dimensions, while `Tensor[(1, 1, 256), float32]` has 3 dimensions; unable to unify: `Tensor[(256), float32]` and `Tensor[(1, 1, 256), float32]`; ;
%542 = nn.max_pool2d(%541, pool_size=[2, 2], strides=[2, 2], layout="NHWC");
%543 = nn.conv2d(%541, %v_param_428, strides=[2, 2], channels=128, kernel_size=[2, 2], data_layout="NHWC", kernel_layout="HWIO");
%544 = nn.bias_add(%543, %v_param_429, axis=3);
%545 = nn.prelu(%544, %v_param_430, axis=3) tensor type `Tensor[(128), float32]` has 1 dimensions, while `Tensor[(1, 1, 128), float32]` has 3 dimensions; unable to unify: `Tensor[(128), float32]` and `Tensor[(1, 1, 128), float32]`; ;
%546 = nn.pad(%545, pad_width=[[0, 0], [1, 1], [1, 1], [0, 0]]);

 %620 = nn.bias_add(%619, %v_param_490, axis=3);
%621 = add(%612, %620);
%622 = nn.prelu(%621, %v_param_491, axis=3) tensor type `Tensor[(256), float32]` has 1 dimensions, while `Tensor[(1, 1, 256), float32]` has 3 dimensions; unable to unify: `Tensor[(256), float32]` and `Tensor[(1, 1, 256), float32]`; ;
%623 = nn.conv2d(%622, %v_param_492, channels=128, kernel_size=[1, 1], data_layout="NHWC", kernel_layout="HWIO");
%624 = nn.bias_add(%623, %v_param_493, axis=3);
%625 = nn.prelu(%624, %v_param_494, axis=3) tensor type `Tensor[(128), float32]` has 1 dimensions, while `Tensor[(1, 1, 128), float32]` has 3 dimensions; unable to unify: `Tensor[(128), float32]` and `Tensor[(1, 1, 128), float32]`; ;
%626 = nn.pad(%625, pad_width=[[0, 0], [1, 1], [1, 1], [0, 0]]);
%627 = nn.conv2d(%626, %v_param_495, groups=128, channels=128, kernel_size=[3, 3], data_layout="NHWC", kernel_layout="HWOI");
%628 = nn.bias_add(%627, %v_param_496, axis=3);
%629 = nn.conv2d(%628, %v_param_497, channels=256, kernel_size=[1, 1], data_layout="NHWC", kernel_layout="HWIO");
%630 = nn.bias_add(%629, %v_param_498, axis=3);
%631 = add(%622, %630);
%632 = nn.prelu(%631, %v_param_499, axis=3) tensor type `Tensor[(256), float32]` has 1 dimensions, while `Tensor[(1, 1, 256), float32]` has 3 dimensions; unable to unify: `Tensor[(256), float32]` and `Tensor[(1, 1, 256), float32]`; ;
%633 = nn.conv2d(%632, %v_param_502, channels=42, kernel_size=[2, 2], data_layout="NHWC", kernel_layout="HWIO");

https://www.dropbox.com/s/lr9wmv0dminvd10/Screenshot%202019-11-12%2013.53.12.png?dl=0

https://www.dropbox.com/s/6jppn19hae2yzte/Screenshot%202019-11-12%2013.54.17.png?dl=0

# Tensors
index name type shape buffer quantization
0 input_1 FLOAT32 [1, 256, 256, 3] 0 None
1 conv2d/Kernel FLOAT32 [16, 3, 3, 3] 1 None
2 conv2d/Bias FLOAT32 [16] 2 None
3 conv2d FLOAT32 [1, 128, 128, 16] 0 None
4 p_re_lu/Alpha FLOAT32 [1, 1, 16] 3 None
5 p_re_lu FLOAT32 [1, 128, 128, 16] 0 None
6 conv2d_1/Kernel FLOAT32 [8, 1, 1, 16] 4 None
7 conv2d_1/Bias FLOAT32 [8] 5 None
8 conv2d_1 FLOAT32 [1, 128, 128, 8] 0 None
9 p_re_lu_1/Alpha FLOAT32 [1, 1, 8] 6 None
10 p_re_lu_1 FLOAT32 [1, 128, 128, 8] 0 None
11 depthwise_conv2d/Kernel FLOAT32 [1, 3, 3, 8] 7 None

[apivovarov]

PRELU docs says that H and W can be "shared". That is what we see in my custom tflite model.
PRELU Alpha tensor with index 4 has shape [1,1,16]. It means that H and W are 1 - H and W are shared, alpha parameters only exist for color dimension.
Also, Alpha tensor is 3D, not 4D. It is because Batch dimension is probably always shared and simply omitted in Alpha tensor shape. If TVM Relay needs PRELU Alpha tensor dimension size to be the same as main input tensor dimension size then we should add 1 to Alpha tensor shape for batch dimension. [1,1,16] -> [1,1,1,16]

[FrozenGene]

could you share your custom tflite model of prelu? Seems that we should reshape alpha tensor to 4D make share meet TVM requirement.

[apivovarov]

I can not share the file unfortunately. I'll try to create similar tflite file to reproduce the issue.

[FrozenGene]

Ok. You could refer my unittest and see how the prelu is constructed, because tf doesn’t have prelu, tflite just recognize op pattern to produce prelu. Then we could see how to solve your model problem.

[apivovarov]

Will test it. BTW the model is hand_landmark.tflite. It is publicly available https://github.com/google/mediapipe/blob/master/mediapipe/models/hand_landmark.tflite

[apivovarov]

I tried the fix with hand_landmark.tflite model. Compilation works. The outputs are the same as when using TFLite runtime. TVM performance is 2.5x faster on cpu. Thank you!