Adding support to ExceptionHandling (Pointcept#292)

* Update events.py Added ExceptionHandler class. The purpose is to save the error traceback to the log_file. * Update train.py Integrate ExceptionHandler Context Manager into the training pipeline. * Imported necessar libraries * Rename ExceptionHandler to ExceptionWriter --------- Co-authored-by: Xiaoyang Wu <[email protected]>
lncopilot · Jul 6, 2024 · 7b37078 · 7b37078
1 parent 9a3f603
commit 7b37078
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 2 deletions.
diff --git a/pointcept/engines/train.py b/pointcept/engines/train.py
@@ -27,7  27,7 @@
 from pointcept.utils.logger import get_root_logger
 from pointcept.utils.optimizer import build_optimizer
 from pointcept.utils.scheduler import build_scheduler
-from pointcept.utils.events import EventStorage
 from pointcept.utils.events import EventStorage, ExceptionWriter
 from pointcept.utils.registry import Registry
 
 
@@ -145,7  145,7 @@ def __init__(self, cfg):
         self.register_hooks(self.cfg.hooks)
 
     def train(self):
-        with EventStorage() as self.storage:
         with EventStorage() as self.storage, ExceptionWriter():
             # => before train
             self.before_train()
             self.logger.info(">>>>>>>>>>>>>>>> Start Training >>>>>>>>>>>>>>>>")

diff --git a/pointcept/utils/events.py b/pointcept/utils/events.py
@@ -14,6  14,8 @@
 import time
 import torch
 import numpy as np
 import traceback
 import sys
 
 from typing import List, Optional, Tuple
 from collections import defaultdict
@@ -25,6  27,7 @@
     "TensorboardXWriter",
     "CommonMetricPrinter",
     "EventStorage",
     "ExceptionWriter",
 ]
 
 _CURRENT_STORAGE_STACK = []
@@ -591,3  863,19 @@ def values(self) -> List[Tuple[float, float]]:
             list[(number, iteration)]: content of the current buffer.
         """
         return self._data
 
 
 class ExceptionWriter:
 
     def __init__(self):
         self.logger = logging.getLogger(__name__)
 
     def __enter__(self):
         pass
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         if exc_type:
             tb = traceback.format_exception(exc_type, exc_val, exc_tb)
             formatted_tb_str = "".join(tb)
             self.logger.error(formatted_tb_str)
             sys.exit(1)  # This prevents double logging the error to the console