如何获取TensorFlow Estimator中model_fn的输入形状以设置权重？

嘿，这个问题我之前在做Estimator构建自编码器时也踩过坑！在Estimator的model_fn里处理动态输入形状确实有点绕，不过咱们分两种情况来解决，保证你能搞定：

1. 最省心的方案：用高层API自动处理权重形状 #

其实完全不用手动纠结tf.get_variable的shape参数，TensorFlow的tf.layers.dense会自动根据输入张量的形状来创建对应的权重，代码简洁还不容易出错。比如你的堆叠自编码器可以这么写：

def model_fn(features, labels, mode):
    # 编码器部分：堆叠全连接层
    encoder_h1 = tf.layers.dense(
        features, 
        units=256, 
        activation=tf.nn.relu,
        kernel_initializer=tf.contrib.layers.variance_scaling_initializer(),
        name="encoder_h1"
    )
    encoder_h2 = tf.layers.dense(
        encoder_h1, 
        units=128, 
        activation=tf.nn.relu,
        kernel_initializer=tf.contrib.layers.variance_scaling_initializer(),
        name="encoder_h2"
    )
    
    # 解码器部分：对称堆叠全连接层
    decoder_h1 = tf.layers.dense(
        encoder_h2, 
        units=256, 
        activation=tf.nn.relu,
        kernel_initializer=tf.contrib.layers.variance_scaling_initializer(),
        name="decoder_h1"
    )
    # 输出层要和输入特征维度一致，用tf.shape动态获取
    output = tf.layers.dense(
        decoder_h1, 
        units=tf.shape(features)[1], 
        activation=tf.sigmoid,
        name="output"
    )
    
    # 定义训练、评估、预测逻辑
    if mode == tf.estimator.ModeKeys.TRAIN:
        loss = tf.reduce_mean(tf.square(output - features))  # MSE损失
        optimizer = tf.train.AdamOptimizer(learning_rate=0.001)
        train_op = optimizer.minimize(loss, global_step=tf.train.get_global_step())
        return tf.estimator.EstimatorSpec(mode=mode, loss=loss, train_op=train_op)
    
    elif mode == tf.estimator.ModeKeys.EVAL:
        loss = tf.reduce_mean(tf.square(output - features))
        eval_metric_ops = {"mse": tf.metrics.mean_squared_error(labels=features, predictions=output)}
        return tf.estimator.EstimatorSpec(mode=mode, loss=loss, eval_metric_ops=eval_metric_ops)
    
    elif mode == tf.estimator.ModeKeys.PREDICT:
        predictions = {"reconstructed": output}
        return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions)

这种方式下，你完全不用手动定义权重的shape，高层API会帮你搞定动态输入的适配，尤其是最后一层用tf.shape(features)[1]来匹配输入维度，完美解决动态形状的问题。

2. 手动创建权重的方案（如果一定要这么做） #

如果你坚持要手动用tf.get_variable来定义权重，那得先从features里提取输入特征的维度：

• 如果你的输入特征维度是固定的（比如MNIST的784维），可以直接从静态形状中提取：

def model_fn(features, labels, mode):
    # 从静态形状中获取输入维度（返回整数，比如784）
    input_dim = features.get_shape()[1].value
    W_init = tf.contrib.layers.variance_scaling_initializer()
    
    # 手动定义编码器权重
    w_encoder1 = tf.get_variable(
        name="encoder_w1",
        shape=[input_dim, 256],  # 这里用静态维度
        initializer=W_init
    )
    b_encoder1 = tf.get_variable(name="encoder_b1", shape=[256], initializer=tf.zeros_initializer())
    encoder_h1 = tf.nn.relu(tf.matmul(features, w_encoder1) + b_encoder1)
    
    # 后续层以此类推...

• 如果输入维度是动态变化的（比较少见，但也能处理），就得用动态形状配合tf.Variable来初始化：

def model_fn(features, labels, mode):
    input_dim = tf.shape(features)[1]
    hidden_dim = 256
    W_init = tf.contrib.layers.variance_scaling_initializer()
    
    # 动态初始化权重，用tf.Variable而不是tf.get_variable
    w_encoder1 = tf.Variable(W_init([input_dim, hidden_dim]), name="encoder_w1")
    b_encoder1 = tf.Variable(tf.zeros([hidden_dim]), name="encoder_b1")
    encoder_h1 = tf.nn.relu(tf.matmul(features, w_encoder1) + b_encoder1)
    
    # ...后续逻辑

不过这种动态初始化的方式在Estimator里要注意兼容性，因为Estimator的变量初始化逻辑依赖于图构建时的静态形状，所以除非特殊需求，还是优先用第一种高层API的方案。

总之，用tf.layers.dense是最省心的选择，既不用纠结动态形状，代码可读性也高，完全满足堆叠自编码器的需求。

内容的提问来源于stack exchange，提问作者Philip616