Add different image channels warning to export onnx Megvii-BaseDetect…

…ion#1692
vossr · Apr 15, 2024 · 0e314d8 · 0e314d8
1 parent 85c22a4
commit 0e314d8
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/tools/export_onnx.py b/tools/export_onnx.py
@@ -65,6 +65,9 @@ def main():
     exp = get_exp(args.exp_file, args.name)
     exp.merge(args.opts)
 
+    if exp.color_channel != args.num_channels:
+        logger.warning(f"When exporting to ONNX, it is important to ensure that the image channels match those used during training.")
+
     if not args.experiment_name:
         args.experiment_name = exp.exp_name