tensorflow自定义估计器的简单示例

import tensorflow as tf

# 定义自定义估计器
class CustomEstimator(tf.estimator.Estimator):
    def __init__(self):
        super(CustomEstimator, self).__init__(
            model_fn=self.model_fn,
            model_dir='./custom_estimator_model'
        )

    def model_fn(self, features, labels, mode):
        # 定义模型结构
        inputs = tf.feature_column.input_layer(features, feature_columns)
        predictions = tf.layers.dense(inputs, 1)

        if mode == tf.estimator.ModeKeys.PREDICT:
            return tf.estimator.EstimatorSpec(mode, predictions=predictions)

        # 定义损失函数和优化器
        loss = tf.losses.mean_squared_error(labels, predictions)
        optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.1)
        train_op = optimizer.minimize(loss, global_step=tf.train.get_global_step())

        # 定义评估指标
        eval_metric_ops = {
            'mse': tf.metrics.mean_squared_error(labels, predictions)
        }

        return tf.estimator.EstimatorSpec(mode, loss=loss, train_op=train_op, eval_metric_ops=eval_metric_ops)

# 定义特征列
feature_columns = [tf.feature_column.numeric_column('x')]

# 创建自定义估计器实例
estimator = CustomEstimator()

# 定义训练输入函数
train_input_fn = tf.estimator.inputs.numpy_input_fn(
    x={'x': x_train},
    y=y_train,
    batch_size=32,
    num_epochs=None,
    shuffle=True
)

# 训练模型
estimator.train(input_fn=train_input_fn, steps=1000)

# 定义评估输入函数
eval_input_fn = tf.estimator.inputs.numpy_input_fn(
    x={'x': x_eval},
    y=y_eval,
    num_epochs=1,
    shuffle=False
)

# 评估模型
eval_result = estimator.evaluate(input_fn=eval_input_fn)
print(eval_result)