有没有办法腌制自定义TensorFlow.keras度量？

Question

我定义了以下自定义度量标准以在tensorflow中训练我的模型：

import tensorflow as tf
from tensorflow import keras as ks
N_CLASSES = 15

class MulticlassMeanIoU(tf.keras.metrics.MeanIoU):
    def __init__(self,
                 y_true = None,
                 y_pred = None,
                 num_classes = None,
                 name = "Multi_MeanIoU",
                 dtype = None):
        super(MulticlassMeanIoU, self).__init__(num_classes = num_classes,
                                             name = name, dtype = dtype)
        self.__name__ = name

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

    def update_state(self, y_true, y_pred, sample_weight = None):
        y_pred = tf.math.argmax(y_pred, axis = -1)
        return super().update_state(y_true, y_pred, sample_weight)

met = MulticlassMeanIoU(num_classes = N_CLASSES)

训练模型后，我保存模型，我也尝试将自定义对象保存如下：

with open("/some/path/custom_metrics.pkl", "wb") as f:
    pickle.dump(met, f)

但是，当我尝试加载指标时像这样：

with open(path_custom_metrics, "rb") as f:
    met = pickle.load(f)

我总是会遇到一些错误，例如attributeError：'MulticlassMeaniou'对象没有属性'update_state_fn'。

现在，我想知道是否可以腌制一个自定义指标，如果是的话，如何？如果我可以通过模型保存自定义指标，那么它将派上用场，因此，当我将模型加载到另一个Python会话中时，我总是拥有首先加载模型所需的指标。可以通过在加载模型之前将完整代码插入其他脚本来重新定义度量标准，但是，我认为这是不好的样式，如果我会在培训脚本中更改有关指标的问题，并且可能会引起问题。忘记将代码复制到其他脚本。

Answer 1

如果您需要泡菜度量标准，则可能的解决方案是使用__ getState __（）和__ setState __（）方法。在（DE）序列化过程中，如果可用，则调用这两种方法。将这些方法添加到您的代码中，您将拥有所需的东西。我试图使其尽可能一般，以便它适用于任何termric：

    def __getstate__(self):
        variables = {v.name: v.numpy() for v in self.variables}
        state = {
            name: variables[var.name]
            for name, var in self._unconditional_dependency_names.items()
            if isinstance(var, tf.Variable)}
        state['name'] = self.name
        state['num_classes'] = self.num_classes
        return state

    def __setstate__(self, state: Dict[str, Any]):
        self.__init__(name=state.pop('name'), num_classes=state.pop('num_classes'))
        for name, value in state.items():
            self._unconditional_dependency_names[name].assign(value)