Skip to content
This repository was archived by the owner on Dec 16, 2022. It is now read-only.

Commit 30c4271

Browse files
David Fidalgojoelgrus
David Fidalgo
authored andcommitted
Close tensorboard's event files properly at the end of the training (#3085)
* add a tensorboard close call at the end of the training * add tensorboard.close call to the callback trainer * make pylint happy * make sphinx happy (I hope)
1 parent 428c151 commit 30c4271

File tree

3 files changed

+18
-0
lines changed

3 files changed

+18
-0
lines changed

allennlp/training/callbacks/log_to_tensorboard.py

+5
Original file line numberDiff line numberDiff line change
@@ -107,6 +107,11 @@ def epoch_end_logging(self, trainer: 'CallbackTrainer'):
107107
log_to_console=True,
108108
epoch=trainer.epoch_number + 1)
109109

110+
@handle_event(Events.TRAINING_END)
111+
def training_end(self, trainer: 'CallbackTrainer'):
112+
# pylint: disable=unused-argument
113+
self.tensorboard.close()
114+
110115
@classmethod
111116
def from_params(cls, serialization_dir: str, params: Params) -> 'LogToTensorboard': # type: ignore
112117
log_batch_size_period = params.pop_int("log_batch_size_period", None)

allennlp/training/tensorboard_writer.py

+10
Original file line numberDiff line numberDiff line change
@@ -211,3 +211,13 @@ def log_activation_histogram(self, outputs, log_prefix: str) -> None:
211211
else:
212212
# skip it
213213
pass
214+
215+
def close(self) -> None:
216+
"""
217+
Calls the ``close`` method of the ``SummaryWriter`` s which makes sure that pending
218+
scalars are flushed to disk and the tensorboard event files are closed properly.
219+
"""
220+
if self._train_log is not None:
221+
self._train_log.close()
222+
if self._validation_log is not None:
223+
self._validation_log.close()

allennlp/training/trainer.py

+3
Original file line numberDiff line numberDiff line change
@@ -549,6 +549,9 @@ def train(self) -> Dict[str, Any]:
549549

550550
epochs_trained += 1
551551

552+
# make sure pending events are flushed to disk and files are closed properly
553+
self._tensorboard.close()
554+
552555
# Load the best model state before returning
553556
best_model_state = self._checkpointer.best_model_state()
554557
if best_model_state:

0 commit comments

Comments
 (0)