Hi Everyone, I have one model which is one subgraph supported by DPU.
When I try to import Pb frozen model for tvm compilation, partition_for_vitis_ai function successfully determines that it is one subgraph on DPU.
But, when I try same model in onnx format, partition_for_vitis_ai function fails to determine that the model should be one subgraph.
Can anyone explain to me what could be wrong ?
Some debugging steps are provided below
Pb tensorflow: Model functions after partition_for_vitis_ai
{GlobalVar(tvmgen_default_vitis_ai_main_0): fn (%vitis_ai_0_i0: Tensor[(1, 28, 28, 1), float32], Inline=1, Compiler=“vitis_ai”, global_symbol=“tvmgen_default_vitis_ai_main_0”, Primitive=1) → Tensor[(1, 1, 1, 10), float32] {
%0 = nn.conv2d(%vitis_ai_0_i0, meta[relay.Constant][0] /* ty=Tensor[(5, 5, 1, 16), float32] /, strides=[2, 2], padding=[1, 1, 2, 2], channels=16, kernel_size=[5, 5], data_layout=“NHWC”, kernel_layout=“HWIO”) / ty=Tensor[(1, 14, 14, 16), float32] */;
%1 = add(%0, meta[relay.Constant][1] /* ty=Tensor[(1, 1, 1, 16), float32] /) / ty=Tensor[(1, 14, 14, 16), float32] */;
%2 = nn.batch_norm(%1, meta[relay.Constant][2] /* ty=Tensor[(16), float32] /, meta[relay.Constant][3] / ty=Tensor[(16), float32] /, meta[relay.Constant][4] / ty=Tensor[(16), float32] /, meta[relay.Constant][5] / ty=Tensor[(16), float32] /, axis=3, epsilon=0.001f) / ty=(Tensor[(1, 14, 14, 16), float32], Tensor[(16), float32], Tensor[(16), float32]) */;
%3 = %2.0;
%4 = nn.relu(%3) /* ty=Tensor[(1, 14, 14, 16), float32] */;
%5 = nn.conv2d(%4, meta[relay.Constant][6] /* ty=Tensor[(5, 5, 16, 32), float32] /, strides=[2, 2], padding=[1, 1, 2, 2], channels=32, kernel_size=[5, 5], data_layout=“NHWC”, kernel_layout=“HWIO”) / ty=Tensor[(1, 7, 7, 32), float32] */;
%6 = add(%5, meta[relay.Constant][7] /* ty=Tensor[(1, 1, 1, 32), float32] /) / ty=Tensor[(1, 7, 7, 32), float32] */;
%7 = nn.batch_norm(%6, meta[relay.Constant][8] /* ty=Tensor[(32), float32] /, meta[relay.Constant][9] / ty=Tensor[(32), float32] /, meta[relay.Constant][10] / ty=Tensor[(32), float32] /, meta[relay.Constant][11] / ty=Tensor[(32), float32] /, axis=3, epsilon=0.001f) / ty=(Tensor[(1, 7, 7, 32), float32], Tensor[(32), float32], Tensor[(32), float32]) */;
%8 = %7.0;
%9 = nn.relu(%8) /* ty=Tensor[(1, 7, 7, 32), float32] */;
%10 = nn.conv2d(%9, meta[relay.Constant][12] /* ty=Tensor[(3, 3, 32, 64), float32] /, strides=[2, 2], padding=[1, 1, 1, 1], channels=64, kernel_size=[3, 3], data_layout=“NHWC”, kernel_layout=“HWIO”) / ty=Tensor[(1, 4, 4, 64), float32] */;
%11 = add(%10, meta[relay.Constant][13] /* ty=Tensor[(1, 1, 1, 64), float32] /) / ty=Tensor[(1, 4, 4, 64), float32] */;
%12 = nn.batch_norm(%11, meta[relay.Constant][14] /* ty=Tensor[(64), float32] /, meta[relay.Constant][15] / ty=Tensor[(64), float32] /, meta[relay.Constant][16] / ty=Tensor[(64), float32] /, meta[relay.Constant][17] / ty=Tensor[(64), float32] /, axis=3, epsilon=0.001f) / ty=(Tensor[(1, 4, 4, 64), float32], Tensor[(64), float32], Tensor[(64), float32]) */;
%13 = %12.0;
%14 = nn.relu(%13) /* ty=Tensor[(1, 4, 4, 64), float32] */;
%15 = nn.conv2d(%14, meta[relay.Constant][18] /* ty=Tensor[(4, 4, 64, 10), float32] /, strides=[4, 4], padding=[0, 0, 0, 0], channels=10, kernel_size=[4, 4], data_layout=“NHWC”, kernel_layout=“HWIO”) / ty=Tensor[(1, 1, 1, 10), float32] */;
%16 = add(%15, meta[relay.Constant][19] /* ty=Tensor[(1, 1, 1, 10), float32] /) / ty=Tensor[(1, 1, 1, 10), float32] */;
%17 = nn.batch_norm(%16, meta[relay.Constant][20] /* ty=Tensor[(10), float32] /, meta[relay.Constant][21] / ty=Tensor[(10), float32] /, meta[relay.Constant][22] / ty=Tensor[(10), float32] /, meta[relay.Constant][23] / ty=Tensor[(10), float32] /, axis=3, epsilon=0.001f) / ty=(Tensor[(1, 1, 1, 10), float32], Tensor[(10), float32], Tensor[(10), float32]) */; %17.0 }
, GlobalVar(main): fn (%images_in: Tensor[(1, 28, 28, 1), float32]) → Tensor[(1, 10), float32] { %0 = @tvmgen_default_vitis_ai_main_0(%images_in) /* ty=Tensor[(1, 1, 1, 10), float32] /; reshape(%0, newshape=[1, -1]) / flatten/Reshape / / ty=Tensor[(1, 10), float32] */ }}
ONNX: Model functions after partition_for_vitis_ai
{GlobalVar(main): fn (%images_in:0: Tensor[(1, 28, 28, 1), float32]) → Tensor[(1, 10), float32] {
%0 = reshape(%images_in:0, newshape=[-1, 1, 28, 28]) /* ty=Tensor[(1, 1, 28, 28), float32] */;
%1 = layout_transform(%0, src_layout=“NCHW”, dst_layout=“NHWC”) /* ty=Tensor[(1, 28, 28, 1), float32] */;
%2 = @tvmgen_default_vitis_ai_main_0(%1) /* ty=Tensor[(1, 14, 14, 16), float32] */;
%3 = nn.relu(%2) /* ty=Tensor[(1, 14, 14, 16), float32] */;
%4 = layout_transform(%3, src_layout=“NHWC”, dst_layout=“NCHW”) /* ty=Tensor[(1, 16, 14, 14), float32] */;
%5 = nn.conv2d(%4, meta[relay.Constant][0] /* ty=Tensor[(32, 16, 5, 5), float32] /, strides=[2, 2], padding=[1, 1, 2, 2], kernel_size=[5, 5]) / ty=Tensor[(1, 32, 7, 7), float32] */;
%6 = add(%5, meta[relay.Constant][1] /* ty=Tensor[(1, 32, 1, 1), float32] /) / ty=Tensor[(1, 32, 7, 7), float32] */;
%7 = nn.relu(%6) /* ty=Tensor[(1, 32, 7, 7), float32] */;
%8 = nn.conv2d(%7, meta[relay.Constant][2] /* ty=Tensor[(64, 32, 3, 3), float32] /, strides=[2, 2], padding=[1, 1, 1, 1], kernel_size=[3, 3]) / ty=Tensor[(1, 64, 4, 4), float32] */;
%9 = add(%8, meta[relay.Constant][3] /* ty=Tensor[(1, 64, 1, 1), float32] /) / ty=Tensor[(1, 64, 4, 4), float32] */;
%10 = nn.relu(%9) /* ty=Tensor[(1, 64, 4, 4), float32] */;
%11 = nn.conv2d(%10, meta[relay.Constant][4] /* ty=Tensor[(10, 64, 4, 4), float32] /, strides=[4, 4], padding=[0, 0, 0, 0], kernel_size=[4, 4]) / ty=Tensor[(1, 10, 1, 1), float32] */;
%12 = add(%11, meta[relay.Constant][5] /* ty=Tensor[(1, 10, 1, 1), float32] /) / ty=Tensor[(1, 10, 1, 1), float32] */;
%13 = layout_transform(%12, src_layout=“NCHW”, dst_layout=“NHWC”) /* ty=Tensor[(1, 1, 1, 10), float32] */;
%14 = layout_transform(%13, src_layout=“NHWC”, dst_layout=“NCHW”) /* ty=Tensor[(1, 10, 1, 1), float32] */;
%15 = reshape(%14, newshape=[-1, 1, 1, 10]) /* ty=Tensor[(1, 1, 1, 10), float32] /; nn.batch_flatten(%15) / ty=Tensor[(1, 10), float32] */ }
, GlobalVar(tvmgen_default_vitis_ai_main_0): fn (%vitis_ai_0_i0: Tensor[(1, 28, 28, 1), float32], Inline=1, Compiler=“vitis_ai”, global_symbol=“tvmgen_default_vitis_ai_main_0”, Primitive=1) → Tensor[(1, 14, 14, 16), float32] {
%0 = nn.conv2d(%vitis_ai_0_i0, meta[relay.Constant][0] /* ty=Tensor[(5, 5, 1, 16), float32] /, strides=[2, 2], padding=[1, 1, 2, 2], kernel_size=[5, 5], data_layout=“NHWC”, kernel_layout=“HWIO”) / ty=Tensor[(1, 14, 14, 16), float32] /; add(%0, meta[relay.Constant][1] / ty=Tensor[(1, 1, 1, 16), float32] /) / ty=Tensor[(1, 14, 14, 16), float32] */ } }