Unet for TensorRT

I’m trying to accelerate unet with tensorrt with:

engine = trt.utils.uff_to_trt_engine(G_LOGGER, uff_model, parser, 1, 1 << 22)

and i’m getting:

[TensorRT] ERROR: Parameter check failed at: Network.cpp::addConcatenation::152, condition: first->getDimensions().d[j] == dims.d[j] && "All non-channel dimensions must match across tensors."

Unet:

def build_for_tensorrt(input_shape):
    inputs = Input(input_shape, name='inputs')
    
    conv1 = unet_conv_block(inputs, 32)
    pool1 = MaxPooling2D(pool_size=(2, 2))(conv1)

    conv2 = unet_conv_block(pool1, 64)
    pool2 = MaxPooling2D(pool_size=(2, 2))(conv2)

    conv3 = unet_conv_block(pool2, 128)
    pool3 = MaxPooling2D(pool_size=(2, 2))(conv3)

    conv4 = unet_conv_block(pool3, 256)
    pool4 = MaxPooling2D(pool_size=(2, 2))(conv4)
    
    conv5 = unet_conv_block(pool4, 512)
    
    up6 = concatenate([Conv2DTranspose(256, (2, 2), strides=(2, 2), activation='relu', padding='same')(conv5), conv4], axis=3)
    conv6 = unet_conv_block(up6, 256)

    up7 = concatenate([Conv2DTranspose(128, (2, 2), strides=(2, 2), activation='relu', padding='same')(conv6), conv3], axis=3)
    conv7 = unet_conv_block(up7, 128)

    up8 = concatenate([Conv2DTranspose(64, (2, 2), strides=(2, 2), activation='relu', padding='same')(conv7), conv2], axis=3)
    conv8 = unet_conv_block(up8, 64)

    up9 = concatenate([Conv2DTranspose(32, (2, 2), strides=(2, 2), activation='relu', padding='same')(conv8), conv1], axis=3)
    conv9 = unet_conv_block(up9, 32)

    conv10 = Conv2D(1, (1, 1), activation='sigmoid', name='outputs')(conv9)

    model = Model(inputs=inputs, outputs=conv10)

    return model

What the cause of the problem?

Hi, I am trying to convert a tf unet model to uff model, but I met some problem, could you show me how you trans your tf unet model to a uff model?

Conv2DTranspose conversion is buggy. I have an open bug for this. Unfortunately bugs are private so I can’t share it here.