Inputs shape is uknown during training (model subclassing) - tensorflow

Question

Consider following model
class FractalNeuralNetwork(tf.keras.Model):
def __init__(self, class_number):
super(FractalNeuralNetwork, self).__init__()
self.box_counting_patches = [BoxCountingPatch(box_size) for box_size in range(3, 41 + 1, 2)]
self.chebyshev = ChebyshevBinaryPatch()
self.euclidean = EuclideanBinaryPatch()
self.manhattan = ManhattanBinaryPatch()
self.percolation_c = PercolationC()
self.percolation_m = PercolationM()
self.percolation_q = PercolationQ()
self.probability = ProbabilityMatrix()
self.fractal_dimension = FractalDimension()
self.lacunarity = Lacunarity()
self.assemble = AssembleFractalImage()
self.resize = tf.keras.layers.Resizing(width=224, height=224)
self.rescale = tf.keras.layers.Rescaling(scale=1./255)
self.mobilenet_v2 = hub.KerasLayer("https://tfhub.dev/google/tf2-preview/mobilenet_v2/feature_vector/4",
output_shape=[1280],
trainable=False)
self.combine = tf.keras.layers.Add()
self.score = tf.keras.layers.Dense(class_number, activation='softmax')
def call(self, inputs):
inputs = tf.ensure_shape(inputs, self.ensure_input_shape)
box_counting_patches = [box_counting_patch(inputs) for box_counting_patch in self.box_counting_patches]
chebyshev = self.chebyshev(inputs=box_counting_patches)
euclidean = self.euclidean(inputs=box_counting_patches)
manhattan = self.manhattan(inputs=box_counting_patches)
percolation_c = self.percolation_c(inputs=[chebyshev, euclidean, manhattan])
percolation_m = self.percolation_m(inputs=[chebyshev, euclidean, manhattan])
percolation_q = self.percolation_q(inputs=[chebyshev, euclidean, manhattan])
probability = self.probability(inputs=[chebyshev, euclidean, manhattan])
fractal_dimension = self.fractal_dimension(inputs=probability)
lacunarity = self.lacunarity(inputs=probability)
fractal_output = self.assemble(
inputs=[
fractal_dimension,
lacunarity,
percolation_c,
percolation_m,
percolation_q
]
)
fractal_output = self.resize(fractal_output)
fractal_output = self.rescale(fractal_output)
fractal_output = self.mobilenet_v2(fractal_output)
original_output = self.rescale(inputs)
original_output = self.mobilenet_v2(original_output)
combined_output = self.combine([fractal_output, original_output])
output = self.score(combined_output)
return output
Every custom layer here is not trainable, they just perform calculations - extract fractal features from images.
The model is trained with the following code:
model = FractalNeuralNetwork(
class_number=CLASS_NUMBER
)
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
model.fit(
training_set,
validation_data=validation_set,
epochs=1
)
The first batch comes with a normal shape (None, 224, 224, 3) but the second (None, None, None, None) and it break my model. Why does it happen?