Change saved tensorflow model input shape at inference time - tensorflow

Question

I've searched everywhere but couldn't find anything. It looks so weird that nobody have already encountered the same problem as I... Let me explain:
I've trained a Tensorflow 2 custom model. During the training I have used set_shape((None, 320, 320, 14)) so that Tensorflow knows the shape (It couldn't infer it for whatever reason... -_-").
I have also saved my custom model at every 100 epochs using:
model.save(os.path.join('models', 'pb', FLAGS.task_name + '-%i' % epoch))
So for the 100th epoch I will have a folder models/pb/my_name-100 that contains
assets
variables
saved_model.pb
Now, at inference time, I just want to load the model (without all the code). So I have created another piece of code that only loads the model and make a prediction... A basic template looks like:
class NeuralNetwork:
def __init__(self, model):
self.model = tf.keras.models.load_model(model)
def predict(self, input_tensor):
pred = self.model(input_tensor[None, ...])
return pred[0]
Where input_tensor is of size (H, W, 14) and so input_tensor[None, ...] is of size:
(None, H, W, 14).
The problem is that, because I have set the shape during training to be (None, 320, 320, 14)... This stupid Tensorflow expects the input to be (None, 320, 320, 14) -_-"!!!. My Neural Network is a fully convolutional neural network, so I really don't care about the input shape. I set it to be (320, 320, 14) during training for memory reason...
During prediction I'd like to be able to do prediction on any kind of shape.
Obviously, I could do a preprocessing function that extracts patch of size (320, 320) from the input image and tiles them. So for example my input_tensor could be of size (30, 320, 320, 14)
And then after the prediction, I could reconstruct the image from the tiles... But I don't want to do that.
Firstly because It takes a bit of time to create the tiles and reconstruct the image from the tile
Secondly because the result will be a bit off due to 0 padding in the convolution. Which means that I need to create overlapping tiles and average the results on the overlapping part to avoid having artifacts during the reconstruction
So my question is simple:
How can I tell tensorflow to accept any width and height at inference time? Omg it's so bothersome. I can't believe that there are not an easy options available to do that

Answer 1

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.