自定义Estimator训练MNIST GPU报错：softmax_tensor不存在求解

这个报错的核心原因有两个，咱们一步步拆解解决：

原因分析 #

1. 训练模式下Tensor被优化移除：你的softmax_tensor只在predictions字典里定义，但在TRAIN模式下，训练流程只用到了logits和loss，完全没有依赖这个softmax张量。TensorFlow的图优化机制会自动移除没有被任何计算依赖的节点，所以训练时图里根本不存在softmax_tensor。
2. 分布式训练的名字前缀问题：你用了tf.contrib.estimator.replicate_model_fn做GPU分布式训练，每个GPU对应的计算塔会给所有张量加上前缀（比如tower_0/），就算softmax_tensor被保留，它的名字也不是你直接写的softmax_tensor，而是带前缀的版本。

解决方案 #

步骤1：确保训练模式下softmax_tensor被保留 #

修改你的cnn_model_fn，显式创建softmax_tensor并给训练操作添加控制依赖，强制TensorFlow保留这个节点：

def cnn_model_fn(features, labels, mode):
    with tf.variable_scope('my_model', reuse=tf.AUTO_REUSE):
        # ... 前面的层代码保持不变 ...
        
        # Logits Layer
        logits = tf.layers.dense(inputs=dropout, units=10)
        # 显式定义softmax_tensor，单独拎出来方便后续操作
        softmax_tensor = tf.nn.softmax(logits, name="softmax_tensor")
        # 将张量加入集合，方便后续Hook获取
        tf.add_to_collection("softmax_tensor", softmax_tensor)
        predictions = {
            "classes" : tf.argmax(input=logits, axis=1),
            "probabilities" : softmax_tensor
        }
        
        if mode == tf.estimator.ModeKeys.PREDICT:
            return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions)
        
        loss = tf.losses.sparse_softmax_cross_entropy(labels=labels, logits=logits)
        
        if mode == tf.estimator.ModeKeys.TRAIN:
            optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.001)
            optimizer = tf.contrib.estimator.TowerOptimizer(optimizer)
            train_op = optimizer.minimize(loss=loss, global_step=tf.train.get_global_step())
            
            # 添加控制依赖，确保softmax_tensor在训练时被计算，不会被优化掉
            with tf.control_dependencies([softmax_tensor]):
                train_op = tf.identity(train_op)
            
            return tf.estimator.EstimatorSpec(mode=mode, loss=loss, train_op=train_op)
        
        # ... 评估模式代码保持不变 ...

步骤2：解决分布式下的名字匹配问题 #

因为replicate_model_fn会给张量加塔前缀，直接用softmax_tensor这个名字找不到对应的节点，咱们自定义一个Hook来动态获取张量，而不是依赖固定名字：

class SoftmaxLoggingHook(tf.train.SessionRunHook):
    def begin(self):
        # 从集合中获取softmax_tensor
        self.softmax_tensor = tf.get_collection("softmax_tensor")[0]
        self.global_step_tensor = tf.train.get_global_step()
    
    def before_run(self, run_context):
        # 指定要获取的张量和全局步数
        return tf.train.SessionRunArgs({
            "probabilities": self.softmax_tensor,
            "global_step": self.global_step_tensor
        })
    
    def after_run(self, run_context, run_values):
        # 打印指定格式的日志，你可以根据需求调整输出内容
        step = run_values.results["global_step"]
        prob_sample = run_values.results["probabilities"][0]
        print(f"Step {step}: Sample probabilities: {prob_sample}")

最后修改训练代码，用自定义Hook替换原来的LoggingTensorHook：

# 替换原来的logging_hook
mnist_classifier.train(
    input_fn=train_inpf,
    steps=20000,
    hooks=[SoftmaxLoggingHook()]
)

额外提示 #

如果你只是单GPU训练，其实可以不用replicate_model_fn，直接用普通的Estimator定义，这样张量名字不会有前缀，步骤1之后直接用原来的LoggingTensorHook（把tensors_to_log改成{"probabilities": "my_model/softmax_tensor"}，因为你的变量作用域是my_model）也能正常工作。

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