I answer my own question.
Unfortunately, my answer will not satisfy everybody. There are so many convoluted things happening in TF (Not to mention that when you search for help, most of it concern the 1st API... -_-").
Anyway, here is the "solution"
In my Neural Network, I have implemented a custom layer to mimic the pytorch function AdaptiveAvgPool2D. My implementation actually use tf.nn_avg_pool under the hood and need to dynamically compute the kernel_size as well as the stride. Here is my code, for reference:
class AdaptiveAvgPool2d(layers.Layer):
def __init__(self, output_shape, data_format='channels_last'):
super(AdaptiveAvgPool2d, self).__init__(autocast=False)
assert data_format in {'channels_last', 'channels_first'}, \
'data format parameter must be in {channels_last, channels_first}'
if isinstance(output_shape, tuple):
self.out_h = output_shape[0]
self.out_w = output_shape[1]
elif isinstance(output_shape, int):
self.out_h = output_shape
self.out_w = output_shape
else:
raise RuntimeError(f"""output_shape should be an Integer or a Tuple2""")
self.data_format = data_format
def call(self, inputs, mask=None):
# input_shape = tf.shape(inputs)
input_shape = inputs.get_shape().as_list()
if self.data_format == 'channels_last':
h_idx, w_idx = 1, 2
else: # can use else instead of elif due to assert in __init__
h_idx, w_idx = 2, 3
stride_h = input_shape[h_idx] // self.out_h
stride_w = input_shape[w_idx] // self.out_w
k_size_h = stride_h + input_shape[h_idx] % self.out_h
k_size_w = stride_w + input_shape[w_idx] % self.out_w
pool = tf.nn.avg_pool(
inputs,
ksize=[k_size_h, k_size_w],
strides=[stride_h, stride_w],
padding='VALID',
data_format='NHWC' if self.data_format == 'channels_last' else 'NCHW')
return pool
The problem is that, I'm using inputs.get_shape().as_list()
to recover int values and not a Tensor(..., type=int). Indeed, the tf.nn.avg_pool accept a list of Int for both the ksize and the strides parameters...
Put it differently, I couldn't use tf.shape(inputs) because it returns a Tensor(..., type=int) and there is no way to recover a int from a Tensor beside by evaluating it...
The way I have implemented my function worked just fine, the problem is that, Tensorflow infers the size under the hood and save the size of all the tensors inside the .pb file when I save it.
Indeed, you can easily open a .pb file with any TextEditor (SublimeText) and see by ourself the expected TensorShape. In my case it was `TensorShape: [null, 320, 320, 14]
So, using set_shape((None, None, None, 14)) instead of set_shape((None, 320, 320, 14)) or nothing actually doesn't change the problem...
The problem is that the average pooling layer does not accept a dynamic kernel size/strides....
I then realizes that there is a tensorflow function for this actually tfa.layers.AdaptiveAveragePooling2D. So, I might just go with it and it will be fine, right?
Well not exactly, under the hood, this tensorflow function use other tf.function like tf.split. The problem with tf.split is that, if your dimension you want to split is of size X and you want to output a tensor of size Y. If X % Y != 0, when tf.split will throw an error... While Pytorch is much more robust and handle cases were X % Y != 0.
Put it differently, it means that, in order for me to use tfa.layers.AdaptiveAveragePooling2D, I need to be sure that the size of the tensor received by this function is divisible by the scalar I pass to the function.
For example, in my case,
The input image are of size: (320, 320, whatever), the input tensor received by tfa.layers.AdaptiveAveragePooling2D is: (40, 40, whatever).
So it means, that the spatial dimension of my tensor was divided by 8 during training. In order for it to work, I should choose a size that can divide 40. Let's say I choose 5.
It means that during the prediction, my neural network will work if the input dimension that the tfa.layers.AdaptiveAveragePooling2D receives is also divisible by 5. But we already know that my input image is 8x bigger then the tensor receives by tfa.layers.AdaptiveAveragePooling2D, so it means that, at prediction times, I can use whatever image size as long as:
- H % (8*5) == 0 and W % (8 * 5) == 0
Where
H and W are respectively the height and the width of my input image.
To do that, we can just implement a simple function:
new_W = W + W % 40 (40 in this example...)
new_H = H + H % 40 (40 in this example...)
This function will stretch a bit the image but not to much so that it should be just fine.
Summing up:
- My AdaptivePooling uses static shape, but I cannot do otherwise since it uses
tf.nn.avg_pool under the hood that doesn't accept dynamic shape
tfa.layers.AdaptiveAveragePooling2D is a work around, but because it relies on tf.split that is not robust to inexact divide, it is not perfect either- The basic solution is to use
tfa.layers.AdaptiveAveragePooling2D and create a preprocessing function before calling the prediction so that the tensor will work just fine with tf.split constraint
- Finally, this is not a good solution either. Because, during training if I receive a tensor of size
(40, 40) and want an avg output of size (5, 5), it means that I basically average (8, 8) features to retrieve one features.
- The problem is that, if I do that during inference time on a bigger image, I will receive a bigger tensor. Let's say:
(100, 200). But since my output will always be (5, 5), it means that I will, this time, average (20, 40) features to retrieve one feature...
Because of this difference between training and inference if I go with this way of doing, inferring on a bigger image might lead to inconsistent results
In my case, the way to go is to batch the images as I have explained in my first post...
Hope it will help some of you.
