TensorFlow profiling guide

pip install gpu-memory-profiler
pip install tensorflow  # or tensorflow-gpu

pip install tensorflow-metal

from tfmemprof import TensorFlowProfiler
import tensorflow as tf

profiler = TensorFlowProfiler(device='/GPU:0')

# Profile a training step
with profiler.profile_context("training_step"):
    with tf.GradientTape() as tape:
        predictions = model(x_train, training=True)
        loss = loss_fn(y_train, predictions)
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))

# Get results
results = profiler.get_results()
print(f"Peak memory: {results.peak_memory_mb:.2f} MB")
print(f"Duration: {results.duration:.2f}s")

from tfmemprof import MemoryTracker

tracker = MemoryTracker(
    sampling_interval=0.5,      # Sample every 500ms
    alert_threshold_mb=4000,    # Alert at 4GB
    device='/GPU:0',
    enable_logging=True
)

# Add custom alert callback
def handle_alert(alert):
    print(f"Memory alert: {alert['message']}")
    print(f"Current memory: {alert['current_memory_mb']:.2f} MB")
    print(f"Threshold: {alert['threshold_mb']:.2f} MB")

tracker.add_alert_callback(handle_alert)

# Start tracking
tracker.start_tracking()

try:
    # Your training loop
    for epoch in range(epochs):
        for step, (x_batch, y_batch) in enumerate(train_dataset):
            train_step(model, x_batch, y_batch)
            
            # Check current memory
            current_mem = tracker.get_current_memory()
            if step % 100 == 0:
                print(f"Step {step}: {current_mem:.1f} MB")
finally:
    # Stop tracking and get results
    results = tracker.stop_tracking()
    
    print(f"\nTracking Results:")
    print(f"Peak memory: {results.peak_memory:.1f} MB")
    print(f"Average memory: {results.average_memory:.1f} MB")
    print(f"Duration: {results.duration:.1f}s")
    print(f"Samples collected: {len(results.memory_usage)}")
    
    if results.alerts_triggered:
        print(f"Alerts triggered: {len(results.alerts_triggered)}")

from tfmemprof import MemoryAnalyzer

analyzer = MemoryAnalyzer()

# Analyze memory patterns
patterns = analyzer.analyze_memory_patterns(profiling_result)
print(f"Peak memory: {patterns['peak_memory_mb']:.2f} MB")
print(f"Average growth: {patterns['average_growth_mb_per_step']:.2f} MB/step")

# Detect memory leaks
leaks = analyzer.detect_memory_leaks(tracking_result)
if leaks:
    print("\nPotential memory leaks detected:")
    for leak in leaks:
        print(f"  Type: {leak['type']}")
        print(f"  Severity: {leak['severity']}")
        print(f"  Description: {leak['description']}")
        print(f"  Recommendation: {leak['recommendation']}")
else:
    print("No memory leaks detected")

# Get optimization score
optimization = analyzer.score_optimization(profiling_result)
print(f"\nOptimization Score: {optimization['overall_score']:.1f}/10")
print("\nCategory Scores:")
for category, score in optimization['categories'].items():
    print(f"  {category}: {score:.1f}/10")

print("\nTop Recommendations:")
for i, rec in enumerate(optimization['top_recommendations'], 1):
    print(f"  {i}. {rec}")

from tfmemprof import profile_function

@profile_function
@tf.function
def train_step(images, labels):
    with tf.GradientTape() as tape:
        predictions = model(images, training=True)
        loss = loss_object(labels, predictions)
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    return loss

@profile_function(name="validation_step")
@tf.function
def validate_step(images, labels):
    predictions = model(images, training=False)
    v_loss = loss_object(labels, predictions)
    return v_loss

# Functions are automatically profiled
for images, labels in train_dataset:
    loss = train_step(images, labels)

for images, labels in val_dataset:
    val_loss = validate_step(images, labels)

from tfmemprof import MemoryVisualizer
import json

# Export tracking data
tracker.start_tracking()
# ... run your code ...
results = tracker.stop_tracking()

# Save to JSON
with open('tracking_results.json', 'w') as f:
    json.dump({
        'peak_memory': results.peak_memory,
        'average_memory': results.average_memory,
        'duration': results.duration,
        'memory_usage': results.memory_usage,
        'timestamps': results.timestamps,
        'alerts': results.alerts_triggered
    }, f, indent=2)

print("Results saved to tracking_results.json")

# Generate visualizations
visualizer = MemoryVisualizer()

# Memory timeline plot
visualizer.plot_memory_timeline(
    profiling_result,
    save_path='memory_timeline.png'
)

# Memory distribution
visualizer.plot_memory_distribution(
    results.memory_usage,
    save_path='memory_distribution.png'
)

print("Visualizations saved")