Hello, I have a custom keras layer. For simplicity, I input the image directly into the custom layer. Here is my keras code:
input_img = keras.layers.Input( shape=(60, 80, 256 ) ) out = CustomLayer(num_clusters = 16)( input_img ) model = keras.models.Model( inputs=input_img, outputs=out )
Code inside the custom layer:
# K: tf.keras.backend # Shape of self.C is 1x1x1x256x16 print 'x.shape=', x.shape # x.shape=(?, 60, 80, 256) # v = K.expand_dims(x, -1) + self.C #original code <--- v = K.reshape( x, [ K.shape(x), K.shape(x), K.shape(x), K.shape(x), 1 ] ) + self.C print 'v.shape=', v.shape #v.shape=(?, ?, ?, 256, 16) return v
The original code with expand_dims seem to work fine on my desktop with tensorflow/keras. However, as TensorRT doesnot support expand_dims I am bypassing it with reshape.
+self.C to x causes the UFFParser to complain while loading the UFF file.
I was under impression that broadcast should work in this case as noted in https://docs.nvidia.com/deeplearning/sdk/tensorrt-support-matrix/index.html#layers-matrix
[TensorRT] ERROR: UffParser: Parser error: net_vlad_layer_1/add_1: Invalid scale mode, nbWeights: 4096
However, if I do not add self.C to the reshaped x, the UFFParser doesn’t seem to throw the error.
Could help me sort this out please?
C was initialized like this in the keras code:
self.C = self.add_weight( name='cluster_centers', shape=[1,1,1,256, 16], initializer='uniform', trainable=True)