Repeatitively creating tensorflow custom model instance and training inside loop gives error

Question

I have created custom model class using example from tensorflowhere . Then I tried to get custom model summary so I searched here . But there was a problem when I tried to train custom model inside a for loop. First iteration always succeed but following iteration crashes with following error message (Input tensors to a CustomModel must come from tf.keras.Input Received: None)

During debugging I found that in first iteration, super(CustomModel, self).__init__() calls Model.__init__() and doesn't call Functional.__init__() . After that super(CustomModel, self).__init__(inputs=self.input_layer, outputs=self.out) calls Functional.__init__() .

But in second iteration super(CustomModel, self).__init__() calls Functional.__init__() right away.

How can I train this custom model in second iteration? Here is my code:

import os
import numpy as np

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
import tensorflow as tf  # noqa
from tensorflow import keras  # noqa
from enum import Enum, auto  # noqa

BATCH_SIZE = 20


class Algo(Enum):
    linearReg = auto()

class CustomModel(keras.Model):
    def __init__(self, input_shape, algo=Algo.linearReg.value, **kwargs):
        super(CustomModel, self).__init__()
        self.in_shape = input_shape
        self.algo = algo
        # create input layer for model summary
        self.input_layer = keras.layers.Input((self.in_shape[1],))

        self.custom_layer = [keras.layers.experimental.preprocessing.Normalization()]
        if algo == Algo.linearReg.value:
            self.custom_layer.append(keras.layers.Dense(units=1, activation='linear'))

        # get output layer with call method for model summary
        self.out = self.call(self.input_layer)

        # reinitialize with input layer and output
        super(CustomModel, self).__init__(
            inputs=self.input_layer,
            outputs=self.out)

    def call(self, input_tensor, **kwargs):
        x = input_tensor
        for layer in self.custom_layer:
            x = layer(x)
        return x

    def predict(self, x, **kwargs):
        for layer in self.custom_layer:
            x = layer.call(x)
        return x

    def from_config(self, config, custom_objects=None):
        super(CustomModel, self).__init__()

    def get_config(self):
        config = super(CustomModel, self).get_config()
        return config


features, labels = (np.random.sample((100, 2, BATCH_SIZE)), np.random.sample((100, 1, BATCH_SIZE)))
dataset = tf.data.Dataset.from_tensor_slices((features, labels))
x, y = next(iter(dataset))

for i in range(4):
    model = CustomModel(x.shape, Algo.linearReg.value)
    model.summary()
    model.compile(optimizer=keras.optimizers.Ftrl(learning_rate=0.01),
                                                  loss='mse', metrics=['mae'])

    model.fit(dataset, epochs=int(2), verbose=2,
              validation_data=dataset.shuffle(2).take(1))
    y_predict = model.predict(x)

Answer 1

Well analyzing further gives me headache. However the problem seems occur when I reinitialize with input and output. So I figured out alternative solution although I don't like model.model() part.

class CustomModel(keras.Model):
    def __init__(self, input_shape, algo=Algo.linearReg.value, **kwargs):

        super(CustomModel, self).__init__()
        self.in_shape = input_shape
        self.algo = algo
        # create input layer for model summary
        self.input_layer = keras.layers.Input((self.in_shape[1],))

        self.custom_layer = [keras.layers.experimental.preprocessing.Normalization()]
        if algo == Algo.linearReg.value:
            self.custom_layer.append(keras.layers.Dense(units=1, activation='linear'))

        # get output layer with call method for model summary
        self.out = self.call(self.input_layer)

    def model(self):
        return keras.Model(inputs=self.input_layer, 
                           outputs=self.call(self.input_layer))

    def call(self, input_tensor, **kwargs):
        for layer in self.custom_layer:
            input_tensor = layer(input_tensor)
        return input_tensor

    def predict(self, input_tensor, **kwargs):
        for layer in self.custom_layer:
            input_tensor = layer.call(input_tensor)
        return input_tensor

    def from_config(self, config, custom_objects=None):
        super(CustomModel, self).__init__()

    def get_config(self):
        config = super(CustomModel, self).get_config()
        return config


features, labels = (np.random.sample((100, 2, BATCH_SIZE)), np.random.sample((100, 1, 
BATCH_SIZE)))
dataset = tf.data.Dataset.from_tensor_slices((features, labels))
x, y = next(iter(dataset))

for i in range(4):
    model = CustomModel(x.shape, Algo.linearReg.value)
    model.model().summary()
    model.compile(optimizer=keras.optimizers.Ftrl(learning_rate=0.01),
              loss='mse', metrics=['mae'])

    model.fit(dataset, epochs=int(2), verbose=2,
          validation_data=dataset.shuffle(2).take(1))
    y_predict = model.predict(x)