Hi all,
I’m trying to convert some TensorFlow models to TensorRT engines. One model requires a “space_to_depth” operation which is not supported by TensorRT.
This produces the output below and fails with a KeyError exception.
Converting to UFF graph
Warning: No conversion function registered for layer: SpaceToDepth yet.
Converting as custom op SpaceToDepth MyModel/SpaceToDepth
name: "MyModel/SpaceToDepth"
op: "SpaceToDepth"
input: "MyModel/SomeTensor"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
attr {
key: "_output_shapes"
value {
list {
shape {
dim {
size: -1
}
dim {
size: 13
}
dim {
size: 13
}
dim {
size: 256
}
}
}
}
}
attr {
key: "block_size"
value {
i: 2
}
}
attr {
key: "data_format"
value {
s: "NHWC"
}
}
...
File "/usr/lib/python3.5/dist-packages/uff/converters/tensorflow/converter.py", line 189, in convert_tf2uff_field
'type': 'dtype', 'list': 'list'}[code]
KeyError: 'shape'
I saw that another user also tried to use this op and was suggested to use a workaround implementing this op using reshape and transpose operators, which I tried but the reshape that is required wants to increase the number of tensor dimensions which TensorRT also does not seem to like.
def space_to_depth1(x: tf.Tensor, block_size: int) -> tf.Tensor:
'''
https://stackoverflow.com/questions/44357970/how-to-implement-tf-space-to-depth-with-numpy
This one is valid but does not work in TensorRT because the reshape changes the number of dimensions
'''
_, height, width, depth = x.get_shape()
reduced_height = height // block_size
reduced_width = width // block_size
result = tf.reshape(x, (-1, reduced_height, block_size, reduced_width, block_size, depth))
result = tf.transpose(result, (0, 1, 3, 2, 4, 5))
result = tf.reshape(result, (-1, reduced_height, reduced_width, block_size * block_size * depth))
return result
Does anyone have an idea on how to implement this layer such that TensorRT supports it?
Also it would be nice if TensorRT could provide better support for the tf.reshape op because I have had some other layers which I could not implement because of this.