update

Author: robin-p-schmitt

returnn/torch/engine.py

@@ -20,7 +20,6 @@
 from torch.utils.data import DataLoader
 from torch import autocast
 from torch.cuda import amp
-from torch.utils.tensorboard import SummaryWriter
 import numpy as np
 
 import returnn
@@ -132,7 +131,12 @@ def __init__(self, config: Config):
         self._reset_dev_memory_caches = config.bool("reset_dev_memory_caches", False)
         self._forward_auto_split_batch_on_oom = config.bool("forward_auto_split_batch_on_oom", False)
         self._stop_on_nonfinite_train_score = config.bool("stop_on_nonfinite_train_score", True)
-        self._tensorboard_writer = SummaryWriter()
+
+        if config.bool("use_tensorboard", False):
+            from torch.utils.tensorboard import SummaryWriter
+            self._tensorboard_writer = SummaryWriter()
+        else:
+            self._tensorboard_writer = None
 
         default_float_dtype = config.value("default_float_dtype", None)
         if default_float_dtype is not None:
@@ -257,7 +261,8 @@ def train(self):
             self.init_train_epoch()
             self.train_epoch()
 
-        self._tensorboard_writer.close()
+        if self._tensorboard_writer:
+            self._tensorboard_writer.close()
 
         print(f"Finished training at epoch {self.epoch}, global train step {self.global_train_step}", file=log.v3)
 
@@ -485,9 +490,10 @@ def train_epoch(self):
                     batch_size_info=_get_batch_size_info(extern_data) if self._log_batch_size else None,
                     log_memory_usage_device=self._device if self._log_memory_usage else None,
                 )
-                # write losses/errors to tensorboard
-                for key, val in eval_info.items():
-                    self._tensorboard_writer.add_scalar(f"train/{key}", val, global_step=self.global_train_step)
+                if self._tensorboard_writer:
+                    # write losses/errors to tensorboard
+                    for key, val in eval_info.items():
+                        self._tensorboard_writer.add_scalar(f"train/{key}", val, global_step=self.global_train_step)
 
                 if self._stop_on_nonfinite_train_score:
                     if any(np.isinf(v) or np.isnan(v) for v in accumulated_losses_dict.values()):
@@ -672,13 +678,12 @@ def eval_model(self, *, skip_already_evaluated: bool = False):
                         start_elapsed=step_end_time - eval_start_time,
                         log_memory_usage_device=self._device if self._log_memory_usage else None,
                     )
-                    # write losses/errors to tensorboard
-                    for key, val in eval_info.items():
-                        self._tensorboard_writer.add_scalar(
-                            f"{dataset_name}/{key}",
-                            val,
-                            global_step=self.global_train_step
-                        )
+                    if self._tensorboard_writer:
+                        # write losses/errors to tensorboard
+                        for key, val in eval_info.items():
+                            self._tensorboard_writer.add_scalar(
+                                f"{dataset_name}/{key}", val, global_step=self.global_train_step
+                            )
 
                     step_idx += 1