Keras自定义训练过程中实时评估模型的Metrics Class

描述

在深度学习任务中经常需要自定义函数类来满足独特的模型评估，基于Keras的自定义Metrics Class的代码实现如下。

代码实现

以自定义Accuracy为例展示：
如果希望指标曲线更加平滑则可以使用累加方式存储每一个step的指标，对应到代码中即使用assign_add；如果不适用累加方式存储指标，则使用assign。

import tensorflow as tf
from tensorflow.keras.metrics import Metric
import tensorflow.keras.backend as K

class AccuracyAngle(Metric):
    def __init__(self, sparse_matrix, batch_size, img_size, thre, metric_mask, angle_range_split_matrix, angle_interval, **kwargs):
        super(AccuracyAngle, self).__init__()
        self.sparse_matrix = sparse_matrix
        self.batch_size = batch_size
        self.thre = thre
        self.img_size = img_size
        self.metric_mask = metric_mask # batchSize*128*128
        self.angle_interval = angle_interval
        self.angle_range_split_matrix = angle_range_split_matrix # 12*360*batchSize
        self.accuracy_angle = self.add_weight(name='accuracy_angle',initializer='zeros')
        self.true_pred = self.add_weight(name='true_pred',initializer='zeros')
        self.total_element = self.add_weight(name='total_element',initializer='zeros')

    def update_state(self, y_true, y_pred, sample_weight=None):
    	'customed operation'
        '# numpy calculation'
        sparseM = tf.sparse.from_dense(K.constant(self.sparse_matrix))
        '# tensor calculation'
        y_pred_ = tf.squeeze(y_pred)*K.constant(self.metric_mask) # batchSize*128*128
        y_true_ = tf.squeeze(y_true)*K.constant(self.metric_mask) # batchSize*128*128      
        y_pred_ = tf.where(tf.math.greater_equal(y_pred_,self.thre), y_pred_, 0) # batchSize*128*128
        y_true_ = tf.where(tf.math.greater_equal(y_true_,self.thre), y_true_, 0) # batchSize*128*128
        y_pred_ = tf.transpose(y_pred_, perm=(1,2,0)) # 128*128*batchSize
        y_true_ = tf.transpose(y_true_, perm=(1,2,0)) # 128*128*batchSize

        chemo_pred = tf.sparse.sparse_dense_matmul(sparseM, tf.reshape(y_pred_,[self.img_size*self.img_size,self.batch_size])) # 360*batchSize
        chemo_true = tf.sparse.sparse_dense_matmul(sparseM, tf.reshape(y_true_,[self.img_size*self.img_size,self.batch_size])) # 360*batchSize
        full_ones = tf.ones(tf.shape(chemo_pred))
        chemo_pred = tf.where(tf.math.greater(chemo_pred,0), full_ones, chemo_pred) # 360*batchSize
        chemo_true = tf.where(tf.math.greater(chemo_true,0), full_ones, chemo_true) # 360*batchSize
        chemo_pred = tf.reduce_sum(K.constant(self.angle_range_split_matrix)*chemo_pred, axis=1)/self.angle_interval # 12*batchSize
        chemo_true = tf.reduce_sum(K.constant(self.angle_range_split_matrix)*chemo_true, axis=1)/self.angle_interval # 12*batchSize

        'calculate accuracy'
        true_pred = K.sum(tf.where(tf.math.equal(K.round(chemo_true), K.round(chemo_pred)),1.0,0.0))
        total_element = K.sum(K.round(chemo_true))+K.sum(1-K.round(chemo_true))
        # acc = (true_pred) / (total_element + K.epsilon())
        self.true_pred.assign_add(true_pred)
        self.total_element.assign_add(total_element)
        # self.accuracy_angle.assign(acc)
        self.accuracy_angle.assign(self.true_pred/(self.total_element+K.epsilon()))
        
    def result(self):
        self.accuracy_angle

        return self.accuracy_angle

    def reset_state(self):
        # The state of the metric will be reset at the start of each epoch.
        self.accuracy_angle.assign(0.)
        self.true_pred.assign(0.)
        self.total_element.assign(0.)

    def get_config(self):
        config = super().get_config().copy()
        config.update(
            {
                'sparse_matrix':self.sparse_matrix,
                'batch_size':self.batch_size,
                'thre':self.thre,
                'img_size':self.img_size,
                'metric_mask':self.metric_mask,
                'angle_interval':self.angle_interval,
                'angle_range_split_matrix':self.angle_range_split_matrix,
                'true_pred':K.eval(self.true_pred),
                'total_element':K.eval(self.total_element),
                'accuracy_angle':K.eval(self.accuracy_angle)
            }
        )
        return config