TFMemoryProfiler
Main TensorFlow GPU memory profiler class for capturing memory snapshots and profiling functions.
Constructor
TFMemoryProfiler(
device: Optional[str] = None,
enable_tensor_tracking: bool = True
)
TensorFlow device name (e.g., ‘/GPU:0’, ‘/CPU:0’). Defaults to auto-detected device.
Whether to track individual tensor lifecycles
Methods
capture_snapshot
Capture current memory state at a point in time.
def capture_snapshot(self, name: str = "snapshot") -> MemorySnapshot
Name identifier for the snapshot
Snapshot containing timestamp, GPU/CPU memory usage, tensor count, and utilization metrics
profile_function
Decorator to profile function memory usage.
def profile_function(self, func: Callable[..., Any]) -> Callable[..., Any]
profiler = TFMemoryProfiler()
@profiler.profile_function
def train_model(model, data):
# Training code
pass
profile_context
Context manager for profiling code blocks.
@contextmanager
def profile_context(self, name: str = "context") -> Iterator[None]
Name for the profiling context
profiler = TFMemoryProfiler()
with profiler.profile_context("training_epoch"):
model.fit(x_train, y_train)
start_continuous_profiling
Start continuous memory profiling in background thread.
def start_continuous_profiling(self, interval: float = 1.0) -> None
Sampling interval in seconds
stop_continuous_profiling
Stop continuous memory profiling.
def stop_continuous_profiling(self) -> None
get_results
Get comprehensive profiling results.
def get_results(self) -> ProfileResult
Average memory usage in MB
Minimum memory usage in MB
Total profiling duration in seconds
Memory growth rate in MB/second
List of captured memory snapshots
Per-function profiling statistics
reset
Reset profiler state and clear all data.
Usage as context manager
The profiler can be used as a context manager:
with TFMemoryProfiler() as profiler:
# Your code here
model.fit(x_train, y_train)
results = profiler.get_results()
MemorySnapshot
Represents a point-in-time memory snapshot.
Unix timestamp of snapshot
Current GPU memory usage in MB
Current CPU memory usage in MB
Reserved GPU memory in MB
GPU memory utilization percentage
ProfileResult
Comprehensive profiling results dataclass.
Total memory allocations detected
Total memory deallocations detected
Properties
duration: Total profiling duration in seconds
memory_growth_rate: Memory growth rate in MB/second
Example
import tensorflow as tf
from tfmemprof import TFMemoryProfiler
# Create profiler
profiler = TFMemoryProfiler(device='/GPU:0')
# Profile a function
@profiler.profile_function
def build_model():
model = tf.keras.Sequential([
tf.keras.layers.Dense(512, activation='relu'),
tf.keras.layers.Dense(10, activation='softmax')
])
return model
# Capture snapshots
profiler.capture_snapshot('before_training')
# Profile code blocks
with profiler.profile_context('training'):
model = build_model()
model.compile(optimizer='adam', loss='categorical_crossentropy')
model.fit(x_train, y_train, epochs=5)
profiler.capture_snapshot('after_training')
# Get results
results = profiler.get_results()
print(f"Peak memory: {results.peak_memory_mb:.2f} MB")
print(f"Average memory: {results.average_memory_mb:.2f} MB")
print(f"Memory growth rate: {results.memory_growth_rate:.2f} MB/s")