Dynamic Memory Allocation#

This script demonstrates how to use dynamic memory allocation with Torch-TensorRT to reduce GPU memory footprint. When enabled, TensorRT engines allocate and deallocate resources dynamically during inference, which can significantly reduce peak memory usage.

This is particularly useful when:

Running multiple models on the same GPU
Working with limited GPU memory
Memory usage needs to be minimized between inference calls

Imports and Model Definition#

[ ]:

import gc
import time

import numpy as np
import torch
import torch_tensorrt as torch_trt
import torchvision.models as models

np.random.seed(5)
torch.manual_seed(5)
inputs = [torch.rand((100, 3, 224, 224)).to("cuda")]

Compilation Settings with Dynamic Memory Allocation#

Key settings for dynamic memory allocation:

dynamically_allocate_resources=True: Enables dynamic resource allocation
lazy_engine_init=True: Delays engine initialization until first inference
immutable_weights=False: Allows weight refitting if needed

With these settings, the engine will allocate GPU memory only when needed and deallocate it after inference completes.

[ ]:

settings = {
    "ir": "dynamo",
    "use_python_runtime": False,
    "enabled_precisions": {torch.float32},
    "immutable_weights": False,
    "lazy_engine_init": True,
    "dynamically_allocate_resources": True,
}

model = models.resnet152(pretrained=True).eval().to("cuda")
compiled_module = torch_trt.compile(model, inputs=inputs, **settings)
print((torch.cuda.mem_get_info()[1] - torch.cuda.mem_get_info()[0]) / 1024**3)
compiled_module(*inputs)

Runtime Resource Allocation Control#

You can control resource allocation behavior at runtime using the ResourceAllocationStrategy context manager. This allows you to:

Switch between dynamic and static allocation modes
Control when resources are allocated and deallocated
Optimize memory usage for specific inference patterns

In this example, we temporarily disable dynamic allocation to keep resources allocated between inference calls, which can improve performance when running multiple consecutive inferences.

[ ]:

time.sleep(30)
with torch_trt.dynamo.runtime.ResourceAllocationStrategy(
    compiled_module, dynamically_allocate_resources=False
):
    print(
        "Memory used (GB):",
        (torch.cuda.mem_get_info()[1] - torch.cuda.mem_get_info()[0]) / 1024**3,
    )
    compiled_module(*inputs)
    gc.collect()
    torch.cuda.empty_cache()
    time.sleep(30)
    print(
        "Memory used (GB):",
        (torch.cuda.mem_get_info()[1] - torch.cuda.mem_get_info()[0]) / 1024**3,
    )
    compiled_module(*inputs)

Memory Usage Comparison#

Dynamic memory allocation trades off some performance for reduced memory footprint:

Benefits:

Lower peak GPU memory usage
Reduced memory pressure on shared GPUs

Considerations:

Slight overhead from allocation/deallocation
Best suited for scenarios where memory is constrained
May not be necessary for single-model deployments with ample memory