Reference Counting & Garbage Collection

Overview

Python uses reference counting as its primary memory management strategy: - Reference counting: When reference count reaches 0, object is immediately freed - Garbage collection: Breaks cycles that reference counting can't handle - Memory overhead: Every object has ref count (8+ bytes overhead) - Implications for ML: Long-running processes must avoid memory leaks

Understanding memory management is essential for building production inference servers.

---

Reference Counting Basics

How Reference Counting Works

import sys

# Create an object
x = []  # Reference count = 1
print(sys.getrefcount(x))  # 2 (one from variable x, one from function argument)

# Each assignment creates a new reference
y = x
print(sys.getrefcount(x))  # 3 (variables x, y, plus function argument)

# Delete reference
del y
print(sys.getrefcount(x))  # 2 (back to x and function argument)

# When last reference is deleted, object is immediately freed
del x
# Object is now garbage collected (ref count = 0)

Reference Counting in Action

import sys
import torch

# PyTorch tensor (has ref count)
tensor = torch.randn(1000, 1000)
print(sys.getrefcount(tensor))  # Probably 2 (tensor and function arg)

# Create references
t1 = tensor
t2 = tensor
print(sys.getrefcount(tensor))  # 4 (tensor, t1, t2, function arg)

# Remove references one by one
del t1
print(sys.getrefcount(tensor))  # 3

del t2
print(sys.getrefcount(tensor))  # 2

del tensor
# Tensor freed immediately - GPU memory released right away!

---

Memory Overhead

Python Object Structure

Every Python object has overhead:

import sys

# Measure object overhead
empty_list = []
print(sys.getsizeof(empty_list))  # ~56 bytes (empty list overhead)

# Add one element
empty_list.append(1)
print(sys.getsizeof(empty_list))  # ~88 bytes (list with capacity)

# The actual integer
x = 1
print(sys.getsizeof(x))  # ~28 bytes (integer object overhead)

# List of integers (all have overhead)
int_list = [1, 2, 3, 4, 5]
total_overhead = 0
total_overhead += sys.getsizeof(int_list)  # List overhead
for item in int_list:
    total_overhead += sys.getsizeof(item)  # Each integer overhead

print(total_overhead)  # ~200+ bytes for 5 integers!

# Compare to NumPy array (no per-element overhead)
import numpy as np
arr = np.array([1, 2, 3, 4, 5], dtype=np.int64)
print(sys.getsizeof(arr))  # ~120 bytes for entire array (8 bytes per element, no overhead per element)

Overhead Breakdown

# PyTorch tensor overhead
tensor = torch.randn(1000, 1000)  # 1M float32 values

# Memory used
memory_bytes = sys.getsizeof(tensor)
print(f"Tensor object overhead: {memory_bytes} bytes")

# Actual data
actual_data = 1000 * 1000 * 4  # 1M * 4 bytes per float32
print(f"Actual data: {actual_data / 1e6:.2f}MB")

# Ratio
print(f"Overhead: {memory_bytes / actual_data:.2%}")  # Usually <1% for large tensors

# But for small tensors, overhead is significant
small_tensor = torch.randn(10)  # 10 float32 values
data_bytes = 10 * 4
overhead = sys.getsizeof(small_tensor)
print(f"Small tensor overhead: {overhead / data_bytes:.1f}x")  # 10-100x overhead

---

Garbage Collection: Breaking Cycles

Reference Cycles

Some objects can reference each other, creating cycles that reference counting can't handle:

# Create a cycle
class Node:
    def __init__(self, value):
        self.value = value
        self.next = None

# Create cycle
node1 = Node(1)
node2 = Node(2)

node1.next = node2  # node1 references node2
node2.next = node1  # node2 references node1 - CYCLE!

# Even after deleting both variables, objects stay in memory
# because they reference each other (ref count never reaches 0)
del node1
del node2
# Objects still in memory due to cycle!

# Garbage collector must detect and break the cycle
import gc
gc.collect()  # Now the objects are freed

Garbage Collection Mechanism

import gc

# GC settings
print(f"GC enabled: {gc.isenabled()}")
print(f"GC threshold: {gc.get_threshold()}")  # (700, 10, 10) - run after 700 object allocations

# Manual collection
collected = gc.collect()
print(f"Collected {collected} objects")

# Disable GC (for performance-critical code)
gc.disable()

# Performance-critical inference
for i in range(1000):
    # No GC pauses
    predictions = model(batch)

# Re-enable GC
gc.enable()
gc.collect()

Detecting Cycles

import gc

class DataHolder:
    def __init__(self, data):
        self.data = data
        self.next = None

# Create cycle
holder1 = DataHolder([1, 2, 3])
holder2 = DataHolder([4, 5, 6])
holder1.next = holder2
holder2.next = holder1

# Objects are in GC's tracked objects
gc.collect()

# Find cycles
garbage = gc.garbage  # Objects in cycles after collection
print(f"Garbage objects: {len(garbage)}")

# Inspect cycles
gc.set_debug(gc.DEBUG_SAVEALL)
gc.collect()
for obj in gc.garbage:
    print(type(obj), id(obj))

---

Memory Leaks in ML Systems

Common Leak Pattern: Callbacks

class TrainingCallback:
    def __init__(self):
        self.history = []

    def on_epoch_end(self, metrics):
        self.history.append(metrics)  # Accumulates every epoch!

# In training loop
callback = TrainingCallback()

for epoch in range(1000):
    metrics = train_epoch()
    callback.on_epoch_end(metrics)  # History grows unbounded

print(len(callback.history))  # 1000 entries - okay
print(sys.getsizeof(callback.history))  # Memory keeps growing

# Fix: Use bounded history or circular buffer
from collections import deque

class FixedCallback:
    def __init__(self, max_history=100):
        self.history = deque(maxlen=max_history)

    def on_epoch_end(self, metrics):
        self.history.append(metrics)  # Automatically removes old entries

Leak Pattern: Circular References in Models

import torch
import torch.nn as nn

class BadModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder = nn.Linear(10, 20)
        self.decoder = nn.Linear(20, 10)

        # LEAK: Circular reference!
        self.encoder.parent = self
        self.decoder.parent = self

    def forward(self, x):
        x = self.encoder(x)
        return self.decoder(x)

# When model is deleted, layers still reference parent, parent references layers
model = BadModel()
del model  # Model object NOT freed due to cycle

# Fix: Use WeakRef to avoid cycle
import weakref

class FixedModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder = nn.Linear(10, 20)
        self.decoder = nn.Linear(20, 10)

        # Use weak reference to avoid cycle
        self.encoder.parent_ref = weakref.ref(self)
        self.decoder.parent_ref = weakref.ref(self)

    def get_parent(self):
        return self.decoder.parent_ref()  # Get actual object if it still exists

---

WeakRef for Non-Owning References

When to Use WeakRef

import weakref

# Normal reference (owns the object)
class Registry:
    instances = []

    def __init__(self, name):
        self.name = name
        Registry.instances.append(self)  # Strong reference

    def __del__(self):
        print(f"Deleted: {self.name}")

obj1 = Registry("obj1")
del obj1
# Output: Deleted: obj1
# But obj1 is still in Registry.instances!

# Better: Use weak reference
class WeakRegistry:
    instances = []

    def __init__(self, name):
        self.name = name
        WeakRegistry.instances.append(weakref.ref(self))

    def __del__(self):
        print(f"Deleted: {self.name}")

obj2 = WeakRegistry("obj2")
del obj2
# Output: Deleted: obj2
# obj2 is now garbage collected

# Access weak reference
print(WeakRegistry.instances[0]())  # Returns obj2 if alive, else None

WeakRef for Caches

import weakref

class DataCache:
    def __init__(self):
        self._cache = weakref.WeakValueDictionary()

    def get_data(self, key):
        """Get data from cache or compute."""
        if key in self._cache:
            return self._cache[key]

        # Compute data
        data = self._compute(key)

        # Store in cache (with weak reference)
        self._cache[key] = data
        return data

    def _compute(self, key):
        return torch.randn(100, 100)

cache = DataCache()

# Data is cached
data1 = cache.get_data("key1")
data2 = cache.get_data("key1")  # Retrieved from cache

# When data is no longer referenced, it's automatically removed from cache
del data1
del data2
# Cache entry is automatically removed (weak reference became invalid)

---

Best Practices for ML Systems

Pattern 1: Avoiding Memory Leaks in Long-Running Servers

class InferenceServer:
    def __init__(self, model_path, max_cache_size=1000):
        self.model = torch.load(model_path)
        self.model.eval()

        # Bounded cache to prevent memory leaks
        self.request_cache = {}
        self.max_cache_size = max_cache_size

    def infer(self, request_id, data):
        """Infer and cache result."""
        if request_id in self.request_cache:
            return self.request_cache[request_id]

        # Compute prediction
        with torch.no_grad():
            result = self.model(data)

        # Cache with size limit
        self.request_cache[request_id] = result

        # Remove old entries if cache too large
        if len(self.request_cache) > self.max_cache_size:
            # Remove oldest entry
            oldest_id = next(iter(self.request_cache))
            del self.request_cache[oldest_id]

        return result

# Usage - no memory leak
server = InferenceServer("model.pth")
for i in range(10000):
    result = server.infer(i, batch_data)
    # Cache stays bounded, no memory growth

Pattern 2: Explicit Memory Management

import gc
import torch

class BatchInference:
    def __init__(self, model, batch_size=32):
        self.model = model
        self.batch_size = batch_size

    def infer_batches(self, data_loader, clear_cache_freq=100):
        """Inference with periodic memory cleanup."""
        results = []

        for batch_idx, batch_data in enumerate(data_loader):
            with torch.no_grad():
                batch_output = self.model(batch_data)
            results.append(batch_output)

            # Periodic cleanup
            if (batch_idx + 1) % clear_cache_freq == 0:
                # Flush unused tensors
                torch.cuda.empty_cache()

                # Force garbage collection
                gc.collect()

        return torch.cat(results, dim=0)

Pattern 3: Context Manager for Memory Management

from contextlib import contextmanager
import gc

@contextmanager
def managed_memory():
    """Context manager for memory management."""
    # Disable GC for performance
    gc.disable()
    gc.collect()  # Clean before

    try:
        yield
    finally:
        # Clean after
        gc.collect()
        gc.enable()

# Usage
with managed_memory():
    # Inference without GC pauses
    for batch in data_loader:
        predictions = model(batch)

---

Debugging Memory Issues

Finding Leaks with tracemalloc

import tracemalloc
import torch

tracemalloc.start()

# Code that might leak
for _ in range(1000):
    x = torch.randn(100, 100)
    y = x ** 2

# Get top memory allocations
current, peak = tracemalloc.get_traced_memory()
print(f"Current: {current / 1e6:.1f}MB; Peak: {peak / 1e6:.1f}MB")

# Get detailed snapshot
snapshot = tracemalloc.take_snapshot()
top_stats = snapshot.statistics('lineno')

for stat in top_stats[:3]:
    print(stat)

Monitoring Reference Counts

import sys
import torch

def monitor_refs(obj, name="object"):
    """Monitor reference count over time."""
    ref_count = sys.getrefcount(obj)
    size = sys.getsizeof(obj)
    print(f"{name}: {ref_count} refs, {size} bytes")

# Monitor tensor lifetime
tensor = torch.randn(1000, 1000)
monitor_refs(tensor, "tensor")  # 2

t1 = tensor
monitor_refs(tensor, "tensor (with t1)")  # 3

t2 = tensor
monitor_refs(tensor, "tensor (with t1, t2)")  # 4

del t1
monitor_refs(tensor, "tensor (after del t1)")  # 3

---

Summary: Memory Management Rules

Pattern	Effect	Use Case
Reference counting	Immediate cleanup	Default, efficient
Garbage collection	Cycle detection	Long-running servers
WeakRef	Non-owning reference	Caches, registries
Bounded collections	Prevent growth	Request caches
Explicit cleanup	Force GC	Performance-critical

---

Related Topics

• 02 Memory Layout & Cache Efficiency - Optimizing data layout
• 04 Multithreading Vs Multiprocessing - Memory sharing in processes
• 06 Memory Profiling & Optimization - Profiling tools
• 02 Autograd Implementation - Tensor lifecycle