richardwhitehead
commented
2 months ago Here is a different (and simpler) version of the code. This version works in AutoGraph mode in TensorFlow 2.13.1 on Python 3.8.10, but produces the above error when run in TensorFlow 2.16.1 on Python 3.10.12.
import keras
import tensorflow as tf
import numpy as np
class UnAveragePooling2D(keras.layers.Layer):
def __init__(self, kernel, *, strides, name = 'co1', dtype = None):
super().__init__(trainable=True, name=name, dtype=dtype)
self.kernel = kernel
self.strides = strides
def build(self, input_shape):
pass
def call(self, inputs):
inputs_shape = tf.shape(inputs)
# the averaged image had a minimum size of source.shape * strides, but it could have been larger if its size
# was not a multiple of stride. For now, assume it was a multiple of stride.
dest_shape = (inputs_shape[0], inputs.shape[1] * self.strides, inputs.shape[2] * self.strides, inputs_shape[3])
#dest = tf.zeros(dest_shape)
# width of destination invalid edges that need to be filled, again assuming that source size was a multiple of stride.
dest_start_invalid = (self.strides // 2, self.strides // 2)
dest_end_invalid = (self.strides // 2, self.strides // 2)
strides_f = tf.cast(self.strides, tf.float32)
dest_start_invalid_f = (tf.cast(dest_start_invalid[0], tf.float32), tf.cast(dest_start_invalid[1], tf.float32))
max_source_f = (tf.cast(inputs.shape[1] - 1, tf.float32), tf.cast(inputs.shape[2] - 1, tf.float32))
# biniliear interpolation with distorted edges
dest = tf.zeros((dest_shape[0], 0, dest_shape[2], dest_shape[3]))
for dest_r in range(dest_shape[1]):
dest_r_f = tf.cast(dest_r, tf.float32)
row = tf.zeros((dest_shape[0], 1, 0, dest_shape[3]))
for dest_c in range(dest_shape[2]):
dest_c_f = tf.cast(dest_c, tf.float32)
source_r = self._dest_to_rource(dest_r_f, strides_f, dest_start_invalid_f[0], max_source_f[0])
source_c = self._dest_to_rource(dest_c_f, strides_f, dest_start_invalid_f[1], max_source_f[1])
value = self._bilinear_interpolate(inputs, source_r, source_c)
#dest = self._assign_pixel_batch_values(dest, dest_r, dest_c, value)
value = tf.reshape(value, (inputs_shape[0], 1, 1, inputs_shape[3]))
row = tf.concat([row, value], axis=2)
dest = tf.concat([dest, row], axis=1)
return dest
@tf.function
def _dest_to_rource(self, dest, stride, dest_start_invalid, max_source):
"""Given a destination pixel row or column position, work out the source
pixel location from which the value should be interpolated."""
if dest < dest_start_invalid + stride - 0.5:
return (dest - dest_start_invalid) / (stride - 0.5)
elif dest > dest_start_invalid + (max_source - 1.0) * stride - 0.5:
return ((dest - dest_start_invalid + 0.5) - (max_source - 1.0) * stride) / (stride - 0.5) + max_source - 1.0
else:
return (dest - dest_start_invalid + 0.5) / stride
@tf.function
def _bilinear_interpolate(self, source, r, c):
"""Given a batch of source images, interpolate each from the four pixels surrounding point r,c.
Near the edge, fade to black.
"""
# Algorithm (ignoring edges):
# 1. Round r,c down to get top-right source pixel r0, c0
# 2. Subtract r0, c0 from r, c get 0..1 proportions of a pixel fr, fc
# 3. The four pixels to be sampled are at p00=[r0, c0], p01=[r0,c0+1], p10=[r0+1, c0], and p11=[r0+1, c0+1]
# 4. For each channel, calculate a linear sum of the four pixels, as follows:
# 5. result = p00.(1-fr).(1-fc)
# + p01.(1-fr).fc
# + p10.fr.(1-fc)
# + p11.fr.fc
# To cope with edges and implement fade-to-black: if any of p00, p01, p10 or p11 would be outside the source
# image then use black instead
r0, c0 = tf.cast(tf.floor(r), tf.int32), tf.cast(tf.floor(c), tf.int32)
fr, fc = r - tf.cast(r0, tf.float32), c - tf.cast(c0, tf.float32)
p00 = self._safe_lookup_pixel(source, r0, c0)
p01 = self._safe_lookup_pixel(source, r0, c0+1)
p10 = self._safe_lookup_pixel(source, r0+1, c0)
p11 = self._safe_lookup_pixel(source, r0+1, c0+1)
return p00*(1-fr)*(1-fc) + p01*(1-fr)*fc + p10*fr*(1-fc) + p11*fr*fc
@tf.function
def _safe_lookup_pixel(self, source, r, c):
"""If x,y is a valid index into the source images then return the
batch of pixels at that location, otherwise return a batch of black pixels."""
source_shape = tf.shape(source)
if 0 <= r < source.shape[1] and 0 <= c < source.shape[2]:
return source[:, r, c, :]
else:
return tf.zeros((source_shape[0], source_shape[3]))
TEST_SIZE = 16
data = [list(range(TEST_SIZE))] * TEST_SIZE
data = tf.convert_to_tensor(data, np.float32)
data = tf.reshape(data, (1, TEST_SIZE, TEST_SIZE, 1))
model = keras.models.Sequential()
model.add(keras.Input(shape=(TEST_SIZE, TEST_SIZE, 1), batch_size=1))
model.add(keras.layers.AveragePooling2D((6, 6), padding='same', strides=4))
model.add(UnAveragePooling2D((6, 6), strides=4))
model.compile(run_eagerly=False)
expanded = model.predict(data)
print(expanded[0,:,:,0])
tilakrayal
google-ml-butler[bot]
Issue type
Bug
Have you reproduced the bug with TensorFlow Nightly?
No
Source
binary
TensorFlow version
2.16.1
Custom code
Yes
OS platform and distribution
Linux Ubuntu 20.04
Mobile device
No response
Python version
3.10.12
Bazel version
No response
GCC/compiler version
No response
CUDA/cuDNN version
No response
GPU model and memory
No response
Current behavior?
The following code works and produces the expected output when run in eager mode, but when using AutoGraph it produces bizarre error messages. The line numbers reported are not the actual location of the errors.
Replacing the contents of _safe_lookup_pixel so it always returns zeros makes the error go away, though of course the output is not then correct.
My installation is using docker image based on the NVIDIA image, details below.
Standalone code to reproduce the issue
Dockerfile: