TensorFlow custom model fit

user_7364215 0
Viewed 381

I want to customize TensorFlow model. I need a custom training algorithm like these: I don't want my model to be inside the custom model just the training algorithm.

class CustomModel(keras.Model):
def __init__(self,inputs, outputs, echo=False):
    super().__init__()
    self.echo = echo
def train_step(self, data):
    x, y = data

    with tf.GradientTape() as tape:
        y_pred = self(x, training=True)
        
        loss = self.compiled_loss(y, y_pred, regularization_losses=self.losses)
        print(loss)
        if self.echo:
            print('*')
    trainable_vars = self.trainable_variables
    gradients = tape.gradient(loss, trainable_vars)
    self.optimizer.apply_gradients(zip(gradients, trainable_vars))
    self.compiled_metrics.update_state(y, y_pred)
    return {m.name: m.result() for m in self.metrics}


inputs = keras.Input(shape=(224,224,3))
x = keras.layers.Conv2D(32,(3,3))(inputs)
x = keras.layers.Conv2D(64,3)(x)
x = keras.layers.Conv2D(64,3)(x)
x = keras.layers.AveragePooling2D()(x)
x = keras.layers.Flatten()(x)
x = keras.layers.Dense(64, activation='relu')(x)
x = keras.layers.Dense(3, activation='softmax')(x)
model = CustomModel( inputs, x,echo= True)
model.compile(optimizer="adam", loss="mse", metrics=["mae"])

opt = Adam(learning_rate=0.0001)
model.compile(optimizer=opt, loss='categorical_crossentropy', metrics=['accuracy'])
epochs = 5
history = model.fit_generator(train_generator,
                              validation_data=valid_generator, verbose=1, epochs=epochs)

error:

NotImplementedError: When subclassing the `Model` class, you should implement a `call` method.
keras tensorflow deep-learning python
1 Answers

You don't need to provide them (inputs, outputs) argument in the init function of your sub-class model. You can implement the call method in your sub-class model as follows:

class CustomModel(keras.Model):
  ...
  ...
  # A call function needs to be implemented 
  def call(self, inputs, *args, **kwargs):
      return self(inputs)

update

Based on the comments, here's a possible workaround. You build the model with the provided input/output within init.

class CustomModel(keras.Model):
    def __init__(self, inputs, x, echo=False, **kwargs):#student
        super().__init__(**kwargs)
        self.model = keras.Model(inputs, x)
        self.echo = echo
        
    def call(self, inputs, *args, **kwargs):
        return self.model(inputs)
    
    def train_step(self, data):
        x, y = data

        with tf.GradientTape() as tape:
            y_pred = self.model(x, training=True)  # Forward pass
            ...
            
        # Compute gradients
        trainable_vars = self.model.trainable_variables
        gradients = ...
        # Update weights
        ...
        return {m.name: m.result() for m in self.metrics}
user_92157800
Related