使用tf.estimator

Question

我正在使用tf.estimator.Estimator构建自定义估算器。

tf.estimator.Estimator自动在张量板上记录训练损失。但是，它每批次记录一次。我想显示每个时期的平均损失，我该怎么做？代码段如下：

def FM(features, labels, mode, params):
    # ....network code omitted
    loss = tf.losses.sparse_softmax_cross_entropy(labels=labels, logits=output)
    if mode == tf.estimator.ModeKeys.EVAL:
        return tf.estimator.EstimatorSpec(mode, loss=loss)
    assert mode == tf.estimator.ModeKeys.TRAIN
    optimizer = tf.train.AdamOptimizer(learning_rate=0.01)
    train_op = optimizer.minimize(loss, global_step=tf.train.get_global_step())
    return tf.estimator.EstimatorSpec(mode, loss=loss, train_op=train_op)

def read_dataset(filename, mode, batch_size = 128):
    ds = tf.data.TextLineDataset(filename).skip(1)
    COLUMNS = [...]
    FIELD_DEFAULTS = [...]
    def _parse_line(line):
        fields = tf.decode_csv(line, FIELD_DEFAULTS)
        features = dict(zip(COLUMNS,fields))
        if mode == tf.estimator.ModeKeys.PREDICT:
            return features
        label = features.pop('...')
        return features, label
    dataset = ds.map(_parse_line)

    if mode == tf.estimator.ModeKeys.TRAIN:
        num_epochs = None # loop indefinitely
        dataset = dataset.shuffle(buffer_size=300000)
    else:
        num_epochs = 1 # end-of-input after this
    dataset = dataset.repeat(num_epochs).batch(batch_size)
    return dataset

classifier = tf.estimator.Estimator(
            model_dir=OUTDIR,
            model_fn=FM,
            params={...})

train_spec = tf.estimator.TrainSpec(input_fn=lambda:read_dataset('train_set', tf.estimator.ModeKeys.TRAIN), max_steps=100000)
eval_spec = tf.estimator.EvalSpec(input_fn=lambda:read_dataset('eval_set', tf.estimator.ModeKeys.EVAL))
tf.estimator.train_and_evaluate(
                              classifier,
                              train_spec,
                              eval_spec)