Tensorflow: Custom MaxPooling layer - shape=(None, 64, 28, 28), dtype=float32)) could not be lifted out of a `tf.function`

I'm attempting to create a custom max pooling layer with a stride of 2, pool size of 2, and "same" padding, but I'm getting the following error:

>ValueError: Argument initial_value (Tensor("custom_max_pooling2d/zeros:0", shape=(None, 64, 28, 28), dtype=float32)) could not be lifted out of a tf.function. (Tried to create variable with name='None'). To avoid this error, when constructing tf.Variables inside of tf.function you can create the initial_value tensor in a tf.init_scope or pass a callable initial_value (e.g., tf.Variable(lambda : tf.truncated_normal([10, 40]))). Please file a feature request if this restriction inconveniences you.Call arguments received by layer "custom_max_pooling2d" (type CustomMaxPooling2D):
• inputs=tf.Tensor(shape=(None, 28, 28, 64), dtype=float32)

It's as though the batch size of 128 (specified in the fit method) is being ignored. It's important to me to keep the low level max pooling functionality as opposed to just calling a tensorflow function. The full code is as follows. Everything works except for the CustomMaxPooling2D class.

import tensorflow as tf
fashion_mnist = tf.keras.datasets.fashion_mnist
(x_train_full, y_train_full), (x_test, y_test) = fashion_mnist.load_data()
x_valid, x_train = x_train_full[:5000]/255.0, x_train_full[5000:]/255.0
y_valid, y_train = y_train_full[:5000], y_train_full[5000:]
x_test = x_test/255.0
tensorboard_writer = tf.summary.create_file_writer('tensorboard/')
class ValidationLossCallback(tf.keras.callbacks.Callback):
def on_epoch_end(self, epoch, logs=None):
with tensorboard_writer.as_default():
val_loss = logs['val_loss']
tf.summary.scalar('val_loss', val_loss, step=epoch)
tensorboard_writer.flush()
loss = tf.keras.losses.SparseCategoricalCrossentropy()
opt = tf.keras.optimizers.Adam()
early_stopping = tf.keras.callbacks.EarlyStopping(monitor='val_loss', patience=25, restore_best_weights=True)
class CustomMaxPooling2D(tf.keras.layers.Layer):
def __init__(self):
super(CustomMaxPooling2D, self).__init__()
self.padding_format = tf.constant([[0, 0], [1, 1], [1, 1], [0, 0]])
self.transpose_format = tf.constant([0, 3, 1, 2])
self.reverse_transpose_format = tf.constant([0, 2, 3, 1])
self.batch_size = tf.Variable(128, dtype=tf.int32)
def build(self, input_shape):
_, self.height, self.width, self.channels = input_shape
self.output_width = tf.floor((self.height-1)/2)+1
def call(self, inputs):
self.out= tf.Variable(tf.zeros(shape=(self.batch_size, self.channels, self.width, self.width)))
if self.width%self.output_width!=0:
inputs = tf.pad(tensor=inputs, paddings=self.padding_format, mode='constant')
inputs = tf.transpose(inputs, perm=self.transpose_format)
for b in range(self.batch_size):
for c in range(self.channels):
input_data = inputs[b][c]
for w in range(self.output_width):
for h in range(self.output_width):
self.out[int(b),int(c),int(w),int(h)].assign(tf.reduce_max(input_data[int(w)*2:(int(w)+1)*2, int(h)*2:(int(h)+1)*2]))
return tf.transpose(self.out, perm=self.reverse_transpose_format)
model = tf.keras.models.Sequential([
tf.keras.layers.Conv2D(64,7,activation="relu",padding="same",input_shape=[28,28,1]),
CustomMaxPooling2D(),
tf.keras.layers.Conv2D(128,3,activation="relu",padding="same"),
tf.keras.layers.Conv2D(128,3,activation="relu",padding="same"),
CustomMaxPooling2D(),
tf.keras.layers.Conv2D(256,3,activation="relu",padding="same"),
tf.keras.layers.Conv2D(256,3,activation="relu",padding="same"),
CustomMaxPooling2D(),
tf.keras.layers.Flatten(),
tf.keras.layers.Dense(128, activation="relu"),
tf.keras.layers.Dropout(0.5),
tf.keras.layers.Dense(64, activation="relu"),
tf.keras.layers.Dropout(0.5),
tf.keras.layers.Dense(10, activation="softmax")
])
model.compile(optimizer=opt, loss=loss)
model.fit(x_train, y_train, batch_size=128, epochs=10000,shuffle=True, validation_data=(x_valid, y_valid), callbacks=[early_stopping, ValidationLossCallback()])
predictions = model.predict(x_test)
# Step 5: Evaluate the predictions
accuracy = tf.keras.metrics.Accuracy()
accuracy.update_state(tf.argmax(predictions, axis=1), y_test)
accuracy_result = accuracy.result().numpy()
print("Accuracy: {:.2f}%".format(accuracy_result * 100))
model.save_weights('models/maxpoolmodel.h5')

Any advice/guidance would be appreciated.

Initially I thought it was the batch_size that was having the issue since that's the error output. After running with your code I noticed that it was having an issue with self.out (which uses batch_size as a parameter).

self.out needs to be moved to the build function:

class CustomMaxPooling2D(tf.keras.layers.Layer):
def __init__(self):
super(CustomMaxPooling2D, self).__init__()
self.padding_format = tf.constant([[0, 0], [1, 1], [1, 1], [0, 0]])
self.transpose_format = tf.constant([0, 3, 1, 2])
self.reverse_transpose_format = tf.constant([0, 2, 3, 1])
self.batch_size = tf.Variable(128, dtype=tf.int32)
def build(self, input_shape):
_, self.height, self.width, self.channels = input_shape
self.output_width = tf.floor((self.height-1)/2)+1
self.out= tf.Variable(tf.zeros(shape=(self.batch_size, self.channels, self.width, self.width)))       
def call(self, inputs):
if self.width%self.output_width!=0:
inputs = tf.pad(tensor=inputs, paddings=self.padding_format, mode='constant')
inputs = tf.transpose(inputs, perm=self.transpose_format)
for b in range(self.batch_size):
for c in range(self.channels):
input_data = inputs[b][c]
for w in range(self.output_width):
for h in range(self.output_width):
self.out[int(b),int(c),int(w),int(h)].assign(tf.reduce_max(input_data[int(w)*2:(int(w)+1)*2, int(h)*2:(int(h)+1)*2]))
return tf.transpose(self.out, perm=self.reverse_transpose_format)