使用 custom_metrics 和自定义损失加载 keras model

Question

我通过子类 keras.model 创建了一个 keras model。 我还使用了自定义损失（焦点损失）、自定义指标（子类 keras.metrics）和学习率衰减。 我已经训练了 model 并使用tf.keras.callbacks.ModelCheckpoint(model_path)保存了它。

当我尝试加载 model 时，我收到一条错误消息，显示ValueError: Unable to restore custom object of type _tf_keras_metric currently. Please make sure that the layer implements get_config and from_config when saving. In addition, please use the custom_objects arg when calling load_model() ValueError: Unable to restore custom object of type _tf_keras_metric currently. Please make sure that the layer implements get_config and from_config when saving. In addition, please use the custom_objects arg when calling load_model()

在深入了解错误后，我开始了解如何传递 custom_objects。 然而，在阅读并尝试了一些东西之后，我仍然无法加载 model。有人可以告诉我正确的方法吗？ 我的代码如下：

def get_metrics():
    train_accuracy = tf.keras.metrics.CategoricalAccuracy(name="train_accuracy")
    val_accuracy = tf.keras.metrics.CategoricalAccuracy(name="val_accuracy")
    confusion_matrix = ConfusionMatrixMetric(20)
    return confusion_matrix, train_accuracy, val_accuracy


def loss_fn(labels, logits):

    epsilon = 1e-9
    model_out = tf.nn.softmax(logits, axis=-1) + epsilon
    ce = - (tf.math.log(model_out) * labels)
    weight = labels * tf.math.pow(1 - model_out, gamma)
    fl = alpha * weight * ce
    loss = tf.reduce_max(fl, axis=-1)
    return loss


def get_optimizer(steps_per_epoch, finetune=False):
    lr = 0.001
    if finetune:
        lr = 0.00001
    lr_fn = tf.keras.optimizers.schedules.PiecewiseConstantDecay(
        [steps_per_epoch * 10], [lr, lr / 10], name=None
    )
    opt_op = tf.keras.optimizers.Adam(learning_rate=lr_fn)
    return opt_op
        
class MyModel(keras.Model):
    def compile(self, optimizer, loss_fn, metric_fn):
        super(MyModel, self).compile()
        self.optimizer = optimizer
        self.loss_fn = loss_fn
        self.confusion_matrix, self.train_accuracy, self.val_accuracy = metric_fn()

    def train_step(self, train_data):
        X, y = train_data
        with tf.GradientTape() as tape:
            logits = self(X, training=True)
            loss = self.loss_fn(y, logits)

        # Compute gradients
        trainable_vars = self.trainable_variables
        gradients = tape.gradient(loss, trainable_vars)

        # Update weights
        self.optimizer.apply_gradients(zip(gradients, trainable_vars))

        # compute metrics keeping an moving average
        y_pred = tf.nn.softmax(y, axis=-1)

        self.train_accuracy.update_state(y, y_pred )
        self.confusion_matrix.update_state(y, y_pred)

        update_dict = {"train_accuracy": self.train_accuracy.result()}
        if 'confusion_matrix_metric' in self.metrics_names:
            self.metrics[0].add_results(update_dict)
        return update_dict
    
class ConfusionMatrixMetric(tf.keras.metrics.Metric):
    def __init__(self, num_classes, **kwargs):
        super(ConfusionMatrixMetric, self).__init__(name='confusion_matrix_metric', **kwargs)  # handles base args (e.g., dtype)
        self.num_classes = num_classes
        self.total_cm = self.add_weight("total", shape=(num_classes, num_classes), initializer="zeros")

    def reset_states(self):
        for s in self.variables:
            s.assign(tf.zeros(shape=s.shape))

    def update_state(self, y_true, y_pred, sample_weight=None):
        self.total_cm.assign_add(self.confusion_matrix(y_true, y_pred))
        return self.total_cm

    def result(self):
        return self.process_confusion_matrix()

    def confusion_matrix(self, y_true, y_pred):
        y_pred = tf.math.argmax(y_pred, 1)
        cm = tf.math.confusion_matrix(y_true, y_pred, dtype=tf.float32, num_classes=self.num_classes)
        return cm

    def process_confusion_matrix(self):
        cm = self.total_cm
        diag_part = tf.linalg.diag_part(cm)
        # accuracy = tf.math.reduce_sum(diag_part) / (tf.math.reduce_sum(cm) + tf.constant(1e-15))
        precision = diag_part / (tf.math.reduce_sum(cm, 0) + tf.constant(1e-15))
        recall = diag_part / (tf.math.reduce_sum(cm, 1) + tf.constant(1e-15))
        f1 = 2 * precision * recall / (precision + recall + tf.constant(1e-15))
        return f1

    def add_results(self, output):
        results = self.result()
        for i in range(self.num_classes):
            output['F1_{}'.format(i)] = results[i]

if __name__ == "__main__":
    model_path = 'model/my_custom_model/'
    create_folder(model_path)
    callbacks = [tf.keras.callbacks.ModelCheckpoint(model_path)]
    # train
    model = MyModel(inputs, outputs)
    model.summary()
    opt_op = get_optimizer(100)

    model.compile(optimizer=opt_op,
                  loss_fn=loss_fn,
                  metric_fn=get_metrics)

    model.fit(train_data_gen(),
              epochs=10,
              callbacks=callbacks)

    tf.keras.models.load_model(model_path)

抱歉代码太长。 但只是想确保我所做的一切都是正确且可以理解的。

Answer 1

正如您的错误说明，如果您想要子类化、使用和加载自定义指标，您应该实现get_config方法。

你已经正确地构建了你的度量子类 class tf.Keras.metrics.Metric ，你只需要添加get_config并用它获取你的参数（据我所知，你只有num_classes ）：

def get_config(self):
    base_config = super().get_config()
    return {**base_config, "num_classes": self.num_classes}

此外，加载时，还必须加载自定义指标：

tf.keras.models.load_model(model_path, custom_objects={"ConfusionMatrixMetric": ConfusionMatrixMetric )

不过要注意以下几点（ 摘自 Aurélien Géron 的《动手机器学习与 Scikit-Learn 和 TensorFlow》，第 2 版）：

Keras API 目前只规定了如何使用子类化来定义层模型、回调和正则化器。 如果您构建其他组件（例如损失、指标、初始值设定项或约束）。 使用子类化，它们可能无法移植到其他 Keras 实现。 Keras API 很可能会更新以指定所有这些组件的子类化。