tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 02:24:38 +08:00
Code

Monitor actual gpu memory allocation (#1715)

Browse Source 
monity actual gpu allocation
This commit is contained in:
Ruilong Li(李瑞龙) 2023-04-09 20:47:24 -07:00 committed by GitHub
parent ec603634ed
commit fa72a9767a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 8 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
nerfstudio/engine/trainer.py
View File

@ -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: