tensorflow - 在自定义 tensorflow keras 指标中获取批量大小

Question

我想制作一个自定义指标，如https://www.tensorflow.org/guide/keras/train_and_evaluate#specifying_a_loss_metrics_and_an_optimizer

我的代码看起来像这样

class IOU(tf.keras.metrics.Metric):
def __init__(self, name='iou_part', **kwargs):
    super(IOU, self).__init__(name=name, **kwargs)
    self.iou = self.add_weight(name='iou_part', initializer='zeros')

    self.template_width = 115
    self.template_height = 75

    self.frame_width = 1280
    self.frame_height = 720

    self.corners = tf.constant([[-0.5, 0.1], [-0.5, 0.5], [0.5, 0.5], [0.5, 0.1]], dtype=tf.float32)

    self.epsilon = 1e-6

def update_state(self, y_true, y_pred, sample_weight=None):
    batch_size = y_true.shape[0]
    fake_frame = tf.ones((batch_size, 1, self.frame_height, self.frame_width))
    fake_template = tf.ones((batch_size, 1, self.template_height, self.template_width))

    target = get_perspective_transform(self.corners, tf.reshape(y_true, (-1, 2, 4)))
    output = get_perspective_transform(self.corners, tf.reshape(y_pred, (-1, 2, 4)))

## Compute IOU

但是，这会给出错误“TypeError：Expected int32, got None of type 'NoneType' 相反。” 这是因为执行 model.compile(....) 时 y_true 为 (None, 4, 2) 。将批量大小纳入指标的正确方法是什么？

标签： tensorflowtensorflow2.0tf.keras