refactor: Add get_memory_usage_mb function for memory monitoring

matsengrp · Jun 18, 2024 · 80ffc3b · 80ffc3b
1 parent 30d2ad8
commit 80ffc3b
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/netam/common.py b/netam/common.py
@@ -1,6 +1,7 @@
 import math
 import inspect
 import itertools
+import resource
 import subprocess
 
 import numpy as np
@@ -228,6 +229,12 @@ def print_tensor_devices(scope="local"):
             print(f"{var_name}: {var_value.device}")
 
 
+def get_memory_usage_mb():
+    # Returns the peak memory usage in MB
+    usage = resource.getrusage(resource.RUSAGE_SELF)
+    return usage.ru_maxrss / 1024  # Convert from KB to MB
+
+
 # Reference: https://pytorch.org/tutorials/beginner/transformer_tutorial.html
 class PositionalEncoding(nn.Module):
     def __init__(self, d_model: int, dropout: float = 0.1, max_len: int = 5000):