Преглед изворни кода

fix when INFER_BENCHMARK_WARMUP > actual num of data

gaotingquan пре 1 година
родитељ
комит
5d784d26c0
1 измењених фајлова са 10 додато и 2 уклоњено
  1. 10 2
      paddlex/inference/models/base/basic_predictor.py

+ 10 - 2
paddlex/inference/models/base/basic_predictor.py

@@ -57,9 +57,17 @@ class BasicPredictor(
             self.benchmark.start()
             if INFER_BENCHMARK_WARMUP > 0:
                 output = super().__call__(input)
+                warmup_num = 0
                 for _ in range(INFER_BENCHMARK_WARMUP):
-                    next(output)
-                self.benchmark.warmup_stop(INFER_BENCHMARK_WARMUP)
+                    try:
+                        next(output)
+                        warmup_num += 1
+                    except StopIteration:
+                        logging.warning(
+                            f"There are only {warmup_num} batches in input data, but `INFER_BENCHMARK_WARMUP` has been set to {INFER_BENCHMARK_WARMUP}."
+                        )
+                        break
+                self.benchmark.warmup_stop(warmup_num)
             output = list(super().__call__(input))
             self.benchmark.collect(len(output))
         else: