jim800121chen
1c53253e6e
問題：同一顆含 Softmax 的 tflite，正式站能編出 KL520 真 .nef， 本專案 worker 在 nef 階段撞 UnimplementedFeature [Softmax] exit 6。 根因：runtime ktc 的 eliminate_tail 是 no-op（閹割版、只印警告）， 尾端 Softmax 未被移除；正式站產物實測為 logits 輸出（Softmax 已砍）。 修法（staging 實測驗證、與正式站產物權重段逐 byte 相同）： - onnx/core.py：新增 remove_tail_softmax()，onnx2onnx_flow 後對所有 platform 移除 terminal Softmax（cut_nodes、不用 cut_types 避免誤砍 中間層下游）、多 output rewire 逐項驗證 fail-loud、onnx.checker、 removed_tail_softmax metadata - precheck.py：R3 誤擋收窄——刪 hw_not_support_col 欄位名誤命中、 op capture 改必須、抽不到具體 op 名改放行+warning（真 fail-open）、 刪自由文字猜 op fallback、加噪音 token 過濾 - tests：+16（29 passed；terminal/no-op/多 output/fail-open 全覆蓋） - docs：TDD §12、design-doc ADR-012、PRD 輸出語意變更（機率→logits） 行為變更：所有 platform NEF 輸出改 logits（與正式站一致）、 分類後處理需呼叫端自行補 softmax。 Review：兩輪（0C/3M→0C/0M）；W3 docker integration 6/6 PASS （520+Softmax e2e 產真 .nef 799,844 bytes、720 回歸、precheck 放行）。 Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
271 lines
11 KiB
Python
271 lines
11 KiB
Python
"""Unit tests for tail-Softmax removal in the ONNX worker.
|
||
|
||
這些測試需要 ``onnx``（建小型手工 graph）但**不需要 ktc / toolchain**：
|
||
`remove_tail_softmax` 走的是 repo 內 libs/ONNX_Convertor/optimizer_scripts 的
|
||
``tools.other.remove_nodes``（純 python + onnx），可以在本機直接驗。
|
||
|
||
模型結構模擬 staging 實測的情境（.autoflow/06-testing/reports/
|
||
tflite-520-softmax-removal-verify-2026-07-06.md）：Gemm(logits) → Softmax(terminal)。
|
||
注意：remove_nodes 的 output 重接依賴 value_info 內有 logits 條目（production
|
||
flow 由 onnx2onnx_flow 的 shape inference 補齊），因此手工 graph 要自帶 value_info。
|
||
"""
|
||
|
||
from pathlib import Path
|
||
import sys
|
||
|
||
import pytest
|
||
|
||
onnx = pytest.importorskip("onnx")
|
||
|
||
from onnx import TensorProto, helper # noqa: E402 (after importorskip)
|
||
|
||
ROOT = Path(__file__).resolve().parents[2]
|
||
sys.path.insert(0, str(ROOT))
|
||
|
||
from services.workers.onnx.core import remove_tail_softmax # noqa: E402
|
||
|
||
|
||
def _make_model_with_tail_softmax(softmax_name="tail_softmax"):
|
||
"""input -> Gemm(logits) -> Softmax -> output，Softmax 為 terminal。"""
|
||
gemm = helper.make_node(
|
||
"Gemm", inputs=["x", "w"], outputs=["logits"], name="gemm0"
|
||
)
|
||
softmax = helper.make_node(
|
||
"Softmax", inputs=["logits"], outputs=["probs"], name=softmax_name
|
||
)
|
||
graph = helper.make_graph(
|
||
nodes=[gemm, softmax],
|
||
name="tail_softmax_graph",
|
||
inputs=[
|
||
helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 4]),
|
||
helper.make_tensor_value_info("w", TensorProto.FLOAT, [4, 3]),
|
||
],
|
||
outputs=[helper.make_tensor_value_info("probs", TensorProto.FLOAT, [1, 3])],
|
||
value_info=[
|
||
# logits 的 value_info：remove_nodes 重接 output 時需要（production
|
||
# 由 shape inference 補、這裡手工補）。
|
||
helper.make_tensor_value_info("logits", TensorProto.FLOAT, [1, 3]),
|
||
],
|
||
)
|
||
return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
|
||
|
||
|
||
def _make_model_without_softmax():
|
||
"""input -> Gemm -> Relu -> output，完全沒有 Softmax。"""
|
||
gemm = helper.make_node("Gemm", inputs=["x", "w"], outputs=["logits"], name="gemm0")
|
||
relu = helper.make_node("Relu", inputs=["logits"], outputs=["y"], name="relu0")
|
||
graph = helper.make_graph(
|
||
nodes=[gemm, relu],
|
||
name="no_softmax_graph",
|
||
inputs=[
|
||
helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 4]),
|
||
helper.make_tensor_value_info("w", TensorProto.FLOAT, [4, 3]),
|
||
],
|
||
outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 3])],
|
||
value_info=[
|
||
helper.make_tensor_value_info("logits", TensorProto.FLOAT, [1, 3]),
|
||
],
|
||
)
|
||
return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
|
||
|
||
|
||
def _make_model_with_tail_softmax_missing_logits_value_info():
|
||
"""同 _make_model_with_tail_softmax、但刻意不給 logits 的 value_info。
|
||
|
||
remove_nodes 的 output 重接依賴 value_info；缺失時該 output 會被 Abandon
|
||
→ 單一 output 模型變成 0 個 output → guard 必須 raise。
|
||
"""
|
||
gemm = helper.make_node("Gemm", inputs=["x", "w"], outputs=["logits"], name="gemm0")
|
||
softmax = helper.make_node(
|
||
"Softmax", inputs=["logits"], outputs=["probs"], name="tail_softmax"
|
||
)
|
||
graph = helper.make_graph(
|
||
nodes=[gemm, softmax],
|
||
name="missing_value_info_graph",
|
||
inputs=[
|
||
helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 4]),
|
||
helper.make_tensor_value_info("w", TensorProto.FLOAT, [4, 3]),
|
||
],
|
||
outputs=[helper.make_tensor_value_info("probs", TensorProto.FLOAT, [1, 3])],
|
||
# 刻意不給 value_info
|
||
)
|
||
return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
|
||
|
||
|
||
def _make_multi_output_model_with_tail_softmax(with_logits_value_info=True):
|
||
"""雙 output 模型：
|
||
branch 1: x -> Gemm(logits) -> Softmax -> probs（graph output 1）
|
||
branch 2: x2 -> Relu -> y（graph output 2、與 Softmax 無關）
|
||
"""
|
||
gemm = helper.make_node("Gemm", inputs=["x", "w"], outputs=["logits"], name="gemm0")
|
||
softmax = helper.make_node(
|
||
"Softmax", inputs=["logits"], outputs=["probs"], name="tail_softmax"
|
||
)
|
||
relu = helper.make_node("Relu", inputs=["x2"], outputs=["y"], name="relu0")
|
||
value_info = (
|
||
[helper.make_tensor_value_info("logits", TensorProto.FLOAT, [1, 3])]
|
||
if with_logits_value_info
|
||
else []
|
||
)
|
||
graph = helper.make_graph(
|
||
nodes=[gemm, softmax, relu],
|
||
name="multi_output_graph",
|
||
inputs=[
|
||
helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 4]),
|
||
helper.make_tensor_value_info("w", TensorProto.FLOAT, [4, 3]),
|
||
helper.make_tensor_value_info("x2", TensorProto.FLOAT, [1, 3]),
|
||
],
|
||
outputs=[
|
||
helper.make_tensor_value_info("probs", TensorProto.FLOAT, [1, 3]),
|
||
helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 3]),
|
||
],
|
||
value_info=value_info,
|
||
)
|
||
return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
|
||
|
||
|
||
def _make_model_with_softmax_output_also_consumed():
|
||
"""Softmax 的 output 同時是 graph output、又被另一個節點消費：
|
||
x -> Softmax -> probs（graph output 1）、probs -> Relu -> y（graph output 2）。
|
||
有 children → 非 terminal → 不可移除。
|
||
"""
|
||
softmax = helper.make_node("Softmax", inputs=["x"], outputs=["probs"], name="softmax0")
|
||
relu = helper.make_node("Relu", inputs=["probs"], outputs=["y"], name="relu0")
|
||
graph = helper.make_graph(
|
||
nodes=[softmax, relu],
|
||
name="output_also_consumed_graph",
|
||
inputs=[helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 3])],
|
||
outputs=[
|
||
helper.make_tensor_value_info("probs", TensorProto.FLOAT, [1, 3]),
|
||
helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 3]),
|
||
],
|
||
value_info=[
|
||
helper.make_tensor_value_info("probs", TensorProto.FLOAT, [1, 3]),
|
||
],
|
||
)
|
||
return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
|
||
|
||
|
||
def _make_model_with_mid_graph_softmax():
|
||
"""input -> Softmax -> Relu -> output，Softmax 不是 terminal（有 children）。"""
|
||
softmax = helper.make_node(
|
||
"Softmax", inputs=["x"], outputs=["probs"], name="mid_softmax"
|
||
)
|
||
relu = helper.make_node("Relu", inputs=["probs"], outputs=["y"], name="relu0")
|
||
graph = helper.make_graph(
|
||
nodes=[softmax, relu],
|
||
name="mid_softmax_graph",
|
||
inputs=[helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 3])],
|
||
outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 3])],
|
||
value_info=[
|
||
helper.make_tensor_value_info("probs", TensorProto.FLOAT, [1, 3]),
|
||
],
|
||
)
|
||
return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
|
||
|
||
|
||
def test_remove_tail_softmax_removes_terminal_softmax_and_rewires_logits():
|
||
model = _make_model_with_tail_softmax()
|
||
|
||
removed = remove_tail_softmax(model)
|
||
|
||
assert removed == ["tail_softmax"]
|
||
# Softmax 已不在 graph
|
||
assert all(n.op_type != "Softmax" for n in model.graph.node)
|
||
# graph output 自動重接成 logits（Softmax 的輸入）
|
||
assert [o.name for o in model.graph.output] == ["logits"]
|
||
# 砍完的模型仍是合法 ONNX
|
||
onnx.checker.check_model(model)
|
||
|
||
|
||
def test_remove_tail_softmax_is_noop_when_no_softmax():
|
||
model = _make_model_without_softmax()
|
||
nodes_before = [(n.name, n.op_type) for n in model.graph.node]
|
||
outputs_before = [o.name for o in model.graph.output]
|
||
|
||
removed = remove_tail_softmax(model)
|
||
|
||
assert removed == []
|
||
# graph 完全不變（絕不能壞既有 onnx/bie 輸入的 flow）
|
||
assert [(n.name, n.op_type) for n in model.graph.node] == nodes_before
|
||
assert [o.name for o in model.graph.output] == outputs_before
|
||
|
||
|
||
def test_remove_tail_softmax_keeps_mid_graph_softmax():
|
||
# 非 terminal 的 Softmax 不移除：remove_nodes 是 cut-from-node 語意，
|
||
# 對中間層 Softmax 使用會把下游一併切掉。
|
||
model = _make_model_with_mid_graph_softmax()
|
||
nodes_before = [(n.name, n.op_type) for n in model.graph.node]
|
||
|
||
removed = remove_tail_softmax(model)
|
||
|
||
assert removed == []
|
||
assert [(n.name, n.op_type) for n in model.graph.node] == nodes_before
|
||
assert any(n.op_type == "Softmax" for n in model.graph.node)
|
||
|
||
|
||
def test_remove_tail_softmax_handles_unnamed_softmax_node():
|
||
# 節點沒有名字時要有 fallback 命名（cut_nodes 靠 node.name 比對）。
|
||
model = _make_model_with_tail_softmax(softmax_name="")
|
||
|
||
removed = remove_tail_softmax(model)
|
||
|
||
assert len(removed) == 1
|
||
assert removed[0] # 有給 fallback 名稱
|
||
assert all(n.op_type != "Softmax" for n in model.graph.node)
|
||
assert [o.name for o in model.graph.output] == ["logits"]
|
||
|
||
|
||
def test_remove_tail_softmax_raises_when_all_outputs_dropped():
|
||
# logits 缺 value_info → remove_nodes 把唯一的 output Abandon
|
||
# → graph 剩 0 個 output → 必須 raise、不能存壞模型。
|
||
model = _make_model_with_tail_softmax_missing_logits_value_info()
|
||
|
||
with pytest.raises(RuntimeError, match="no outputs"):
|
||
remove_tail_softmax(model)
|
||
|
||
|
||
def test_remove_tail_softmax_multi_output_rewires_and_keeps_unrelated_output():
|
||
# 多 output 模型：Softmax 餵的 output 重接成 logits、無關的 output 保留。
|
||
model = _make_multi_output_model_with_tail_softmax(with_logits_value_info=True)
|
||
|
||
removed = remove_tail_softmax(model)
|
||
|
||
assert removed == ["tail_softmax"]
|
||
assert all(n.op_type != "Softmax" for n in model.graph.node)
|
||
assert {o.name for o in model.graph.output} == {"logits", "y"}
|
||
onnx.checker.check_model(model)
|
||
|
||
|
||
def test_remove_tail_softmax_multi_output_raises_on_silently_dropped_output():
|
||
# 多 output 模型 + logits 缺 value_info：Softmax 餵的 output 被 Abandon、
|
||
# 但另一個 output 還在（graph.output 非空、checker 也會過）——
|
||
# 逐 output 驗證的 guard 必須抓到並 raise（Major-3）。
|
||
model = _make_multi_output_model_with_tail_softmax(with_logits_value_info=False)
|
||
|
||
with pytest.raises(RuntimeError, match="silently dropped"):
|
||
remove_tail_softmax(model)
|
||
|
||
|
||
def test_remove_tail_softmax_raises_on_duplicate_node_name():
|
||
# cut_nodes 靠名字比對：病態的同名節點（Softmax 與其他節點撞名）
|
||
# 會導致 remove_nodes 誤砍無關節點 → 必須 raise 而非靜默砍。
|
||
model = _make_model_with_tail_softmax(softmax_name="dup")
|
||
model.graph.node[0].name = "dup" # gemm0 改成與 Softmax 同名
|
||
|
||
with pytest.raises(RuntimeError, match="matches 2 nodes"):
|
||
remove_tail_softmax(model)
|
||
|
||
|
||
def test_remove_tail_softmax_keeps_softmax_whose_output_is_also_consumed():
|
||
# Softmax output 同時是 graph output 又被其他節點消費 → 非 terminal、不移除。
|
||
model = _make_model_with_softmax_output_also_consumed()
|
||
nodes_before = [(n.name, n.op_type) for n in model.graph.node]
|
||
outputs_before = [o.name for o in model.graph.output]
|
||
|
||
removed = remove_tail_softmax(model)
|
||
|
||
assert removed == []
|
||
assert [(n.name, n.op_type) for n in model.graph.node] == nodes_before
|
||
assert [o.name for o in model.graph.output] == outputs_before
|