[models] Keep the early stop of MASTER decoding out of the ONNX export

doctr/models/recognition/master/pytorch.py

@@ -253,7 +253,9 @@ def decode(self, encoded: torch.Tensor) -> torch.Tensor:
             output = self.decoder(ys, encoded, source_mask, target_mask)
             # update ys with the next token and ignore the first token (SOS)
             ys[:, i + 1] = self.linear(output[:, i]).argmax(-1)
-            if (ys == self.vocab_size).any(dim=-1).all():  # every sequence has emitted EOS
+            # Stop decoding if all sequences have reached the EOS token
+            # NOTE: `break` isn't correctly translated to Onnx so we don't break here if we want to export
+            if not self.exportable and (ys == self.vocab_size).any(dim=-1).all():
                 break
 
         # Shape (N, max_length, vocab_size + 1)

tests/pytorch/test_models_recognition_pt.py

@@ -300,6 +300,31 @@ def test_models_onnx_export(arch_name, input_shape):
         pytest.skip(f"Output of {arch_name}:\nMax element-wise difference: {np.max(np.abs(pt_logits - ort_outs[0]))}")
 
 
+@pytest.mark.parametrize("arch_name, head", [["master", "linear"], ["parseq", "head"]])
+@pytest.mark.parametrize("exportable", [False, True])
+def test_autoregressive_decoding_early_stop(arch_name, head, exportable):
+    # Decoding stops once every sequence has emitted EOS, except when the model is exportable: tracing would
+    # freeze the number of decoding steps to the one the dummy input needed
+    model = recognition.__dict__[arch_name](pretrained=False, pretrained_backbone=False, exportable=exportable).eval()
+    input_tensor = torch.rand((2, 3, 32, 128))
+
+    def _decoder_calls(eos_bias: float) -> int:
+        calls = []
+        with torch.no_grad():
+            getattr(model, head).bias[model.vocab_size] = eos_bias
+        hook = model.decoder.register_forward_hook(lambda *_: calls.append(1))
+        try:
+            with torch.no_grad():
+                model(input_tensor)
+        finally:
+            hook.remove()
+        return len(calls)
+
+    # EOS emitted at the first step vs. never emitted
+    early, full = _decoder_calls(1e4), _decoder_calls(-1e4)
+    assert early == full if exportable else early < full
+
+
 @pytest.mark.parametrize(
     "arch_name",
     [