TensorFlow 自定义模型拟合

Question

我想自定义 TensorFlow 模型。我需要像这样的自定义训练算法：我不希望我的模型只包含在训练算法中的自定义模型中。

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)

错误：

NotImplementedError: When subclassing the `Model` class, you should implement a `call` method.

Answer 1

您不需要在 sub-class 模型的 init 函数中提供它们（输入、输出）参数。您可以在 sub-class 模型中实现 call 方法，如下所示：

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

更新

根据评论，这是一个可能的解决方法。您在 init.

中使用提供的 input/output 构建模型

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}

TensorFlow 自定义模型拟合

TensorFlow custom model fit

python

deep-learning

keras

tensorflow