Fix jit scripting bugs

zhiqwang · Apr 8, 2021 · 5495d7f · 5495d7f
1 parent 6a06453
commit 5495d7f
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 7 deletions.
diff --git a/test/test_torchscript.py b/test/test_torchscript.py
@@ -1,3 +1,4 @@
+# Copyright (c) 2020, Zhiqiang Wang. All Rights Reserved.
 import unittest
 
 import torch
@@ -65,6 +66,3 @@ def test_yolotr_script(self):
         self.assertTrue(out[0]["scores"].equal(out_script[0]["scores"]))
         self.assertTrue(out[0]["labels"].equal(out_script[0]["labels"]))
         self.assertTrue(out[0]["boxes"].equal(out_script[0]["boxes"]))
-
-if __name__ == "__main__":
-    unittest.main()
diff --git a/yolort/models/yolo.py b/yolort/models/yolo.py
@@ -118,7 +118,7 @@ def forward(
 
         if torch.jit.is_scripting():
             if not self._has_warned:
-                warnings.warn("YOLO always returns a (Losses, Detections) tuple in scripting")
+                warnings.warn("YOLO always returns a (Losses, Detections) tuple in scripting.")
                 self._has_warned = True
             return losses, detections
         else:

diff --git a/yolort/models/yolo_module.py b/yolort/models/yolo_module.py
@@ -1,4 +1,5 @@
 # Copyright (c) 2021, Zhiqiang Wang. All Rights Reserved.
+import warnings
 import argparse
 
 import torch
@@ -8,7 +9,6 @@
 
 from . import yolo
 from .transform import GeneralizedYOLOTransform
-
 from ..datasets import DetectionDataModule, DataPipeline
 
 from typing import Any, List, Dict, Tuple, Optional
@@ -50,6 +50,9 @@ def __init__(
 
         self._data_pipeline = None
 
+        # used only on torchscript mode
+        self._has_warned = False
+
     def forward(
         self,
         inputs: List[Tensor],
@@ -84,13 +87,19 @@ def forward(
 
         if self.training:
             # compute the losses
-            losses = outputs
+            if torch.jit.is_scripting():
+                losses = outputs[0]
+            else:
+                losses = outputs
         else:
             # Rescale coordinate
             detections = self.transform.postprocess(outputs, samples.image_sizes, original_image_sizes)
 
         if torch.jit.is_scripting():
-            return losses, detections
+            if not self._has_warned:
+                warnings.warn("YOLOModule always returns Detections in scripting.")
+                self._has_warned = True
+            return detections
         else:
             return self.eager_outputs(losses, detections)