使用TOCO转换TensorFlow冻结模型为TFLite时维度不匹配错误求助

这个错误的根源很明确——你使用的TFLite版本（从报错时间2018年来看是旧版本）里，ArgMax操作只支持4维输入，且必须在最后一个维度（axis=3）上计算，但你的tf.argmax(input=logits, axis=1)中，logits是2维张量（形状大概是[batch_size, num_classes]），输入维度不符合要求，才触发了input_dims.size() == 4 (2 vs. 4)的断言失败。

具体解决方案 #

你需要修改训练代码中的ArgMax部分，先把2维的logits扩展为4维，在最后一个维度上执行ArgMax，再把结果压缩回原来的维度：

# 替换你原来的predictions定义代码
# 把2维logits扩展为4维：[batch_size, 1, 1, num_classes]
expanded_logits = tf.expand_dims(tf.expand_dims(logits, 1), 1)
# 在最后一个维度（axis=3）执行ArgMax
argmax_4d = tf.argmax(input=expanded_logits, axis=3, name="classes_tensor_4d")
# 把结果压缩回1维（匹配原来的输出形状）
classes_tensor = tf.squeeze(argmax_4d, name="classes_tensor")

predictions = {
    "classes": classes_tensor,
    "probabilities": tf.nn.softmax(logits, name="softmax_tensor")
}

后续操作步骤 #

1. 重新运行训练代码，确保新的节点被正确保存到检查点中；
2. 重新执行冻结模型的流程，生成新的output_graph.pb；
3. 再次用你原来的TOCO命令尝试转换，这时候ArgMax的输入维度就符合要求了。

另外提一句，你终端输出里的IsVariableInitialized、RefSwitch等不支持的操作，如果不影响模型推理，可以暂时忽略——但如果后续还有问题，可能需要在冻结模型时清理掉这些和训练相关的节点（比如只保留推理所需的子图）。

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