From fa72a9767aba00601d0dcd6d4f531fdb49687a35 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ruilong=20Li=28=E6=9D=8E=E7=91=9E=E9=BE=99=29?= Date: Sun, 9 Apr 2023 20:47:24 -0700 Subject: [PATCH] Monitor actual gpu memory allocation (#1715) monity actual gpu allocation --- nerfstudio/engine/trainer.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index 0262703dd..275814492 100644 --- a/nerfstudio/engine/trainer.py +++ b/nerfstudio/engine/trainer.py @@ -238,6 +238,14 @@ class Trainer: writer.put_scalar(name="Train Loss", scalar=loss, step=step) writer.put_dict(name="Train Loss Dict", scalar_dict=loss_dict, step=step) writer.put_dict(name="Train Metrics Dict", scalar_dict=metrics_dict, step=step) + # The actual memory allocated by Pytorch. This is likely less than the amount + # shown in nvidia-smi since some unused memory can be held by the caching + # allocator and some context needs to be created on GPU. See Memory management + # (https://pytorch.org/docs/stable/notes/cuda.html#cuda-memory-management) + # for more details about GPU memory management. + writer.put_scalar( + name="GPU Memory (MB)", scalar=torch.cuda.max_memory_allocated() / (1024**2), step=step + ) # Do not perform evaluation if there are no validation images if self.pipeline.datamanager.eval_dataset: