我从头开始使用Tensorflow实现了一个基本的神经网络，并在MNIST时尚数据集上进行了训练。它训练正确，并且在10个类别上输出测试准确率大约在~88-90%。

现在我编写了一个predict()函数，使用训练好的权重来预测给定图像的类别。以下是代码：

def predict(images, trained_parameters):    Ws, bs = [], []    parameters = {}    for param in trained_parameters.keys():        parameters[param] = tf.convert_to_tensor(trained_parameters[param])    X = tf.placeholder(tf.float32, [images.shape[0], None], name = 'X')    Z_L = forward_propagation(X, trained_parameters)    p = tf.argmax(Z_L) # 正常工作    # p = tf.argmax(tf.nn.softmax(Z_L)) # 如果应用softmax则不工作    with tf.Session() as session:        prediction = session.run(p, feed_dict={X: images})    return prediction

这使用了forward_propagation()函数，该函数返回最后一层的加权和（Z），而不是激活值（A），因为TensorFlow的tf.nn.softmax_cross_entropy_with_logits()需要Z而不是A，因为它会通过应用softmax来计算A。有关详细信息，请参考此链接。

现在在predict()函数中，当我使用Z而不是A（激活值）进行预测时，它工作正常。但是如果我在Z上计算softmax（这是最后一层的激活值A），它会给出错误的预测。

为什么在加权和Z上进行预测是正确的？我们不是应该先应用softmax激活（并计算A），然后再进行预测吗？

如果有人想查看我的完整代码，这里是我的colab笔记本的链接：笔记本Gist链接

那么我在这里遗漏了什么呢？

回答：

大多数TF函数，例如tf.nn.softmax，默认假设批次维度是第一个维度——这是常见的做法。现在，我注意到你的代码中批次维度是第二个，即你的输出形状是(output_dim=10, batch_size=?)，结果，tf.nn.softmax沿着批次维度计算softmax激活值。

不遵循惯例并没有错——人们只需要意识到这一点。沿着第一个轴计算softmax的argmax应该会得到预期的结果（这相当于对logits进行argmax）：

p = tf.argmax(tf.nn.softmax(Z_L, axis=0))

此外，我还建议在网络输入多张图像时，沿着第一个轴计算argmax。