推理时添加softmax

154f42f1 · z37757 · b580fa05 · 154f42f1
隐藏空白更改
内联并排

Showing with 2 addition and 1 deletion

ppocr/modeling/heads/rec_att_head.py ppocr/modeling/heads/rec_att_head.py +2 -1

未找到文件。
--- a/ppocr/modeling/heads/rec_att_head.py
+++ b/ppocr/modeling/heads/rec_att_head.py
@@ -169,7 +169,8 @@ class AttentionLSTM(nn.Layer):
                next_input = probs_step.argmax(axis=1)

                targets = next_input
-
+        if not self.training:
+            probs = paddle.nn.functional.softmax(probs, axis=2)
        return probs