Setting log level of onnx and dynamo to ERROR

Reducing amount of logs that is not relevant to the user
sacdallago · Nov 18, 2024 · 98ee386 · 98ee386
1 parent baf47f3
commit 98ee386
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/biotrainer/solvers/solver.py b/biotrainer/solvers/solver.py
@@ -331,7 +331,7 @@ def save_as_onnx(self, embedding_dimension: int, output_dir: Optional[str] = Non
         self.network.eval()
 
         # Export
-        export_options = torch.onnx.ExportOptions(dynamic_shapes=True)
+        export_options = torch.onnx.ExportOptions(dynamic_shapes=True, op_level_debug=False)
         onnx_program = torch.onnx.dynamo_export(self.network, dummy_input,
                                                 export_options=export_options)
         onnx_file_name = f"{output_dir}/{self.checkpoint_name.split('.')[0]}.onnx"

diff --git a/biotrainer/utilities/executer.py b/biotrainer/utilities/executer.py
@@ -1,3 +1,4 @@
+import torch
 import logging
 
 from ruamel import yaml
@@ -29,6 +30,10 @@ def _setup_logging(output_dir: str):
                             logging.StreamHandler()]
                         )
     logging.captureWarnings(True)
+    # Only log errors for onnx and dynamo
+    torch._logging.set_logs(dynamo=logging.ERROR, onnx=logging.ERROR, onnx_diagnostics=False)
+    for logger_name in ["torch.onnx", "torch._dynamo", "onnxscript"]:
+        logging.getLogger(logger_name).setLevel(logging.ERROR)
 
 
 def _write_output_file(out_filename: str, config: dict) -> None: