HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

Add model.eval() in TensorBoad graph visualization to avoid BN stats changes (#8629)

Browse source 
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Mohammed Yasin 2024-03-06 09:22:28 +08:00 • committed by GitHub
parent 9c42596145
commit 609a0cefbf
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 1 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

1
ultralytics/utils/callbacks/tensorboard.py
View file

@ -45,6 +45,7 @@ def _log_tensorboard_graph(trainer):
# Try simple method first (YOLO) # Try simple method first (YOLO)
with contextlib.suppress(Exception): with contextlib.suppress(Exception):
trainer.model.eval() # place in .eval() mode to avoid BatchNorm statistics changes
WRITER.add_graph(torch.jit.trace(de_parallel(trainer.model), im, strict=False), []) WRITER.add_graph(torch.jit.trace(de_parallel(trainer.model), im, strict=False), [])
LOGGER.info(f"{PREFIX}model graph visualization added ✅") LOGGER.info(f"{PREFIX}model graph visualization added ✅")
return return