tdb - Tensor Debugging Library for PyTorch

tdb is a lightweight, intuitive Python library designed for efficient debugging and exploration of PyTorch tensors. It streamlines development and improves troubleshooting by offering a straightforward interface for inspecting tensor values and shapes, facilitating a better understanding of operations and data flow in PyTorch.

Installation:

pip install git+https://github.com/KonradSzafer/tdb.git

Usage:

Debugging:

import tdb

x = torch.rand(1, 20, 5, 30)
tdb.print(x, 'x')

Memory Management:

Use the following function to free the CUDA memory of the given models and tensors. The function returns objects with their memory released:

import tdb

model, batch, output = tdb.release_memory(model, batch, output)

To see the current memory usage, run:

tdb.print_memory()

Configuration Options:

tdb.options['disable'] = False # Set to True to disable all tdb output
tdb.options['print_values_threshold'] = 10 # Determines the maximum number of values to display from the last dimension of a tensor
tdb.options['assignment_symbol'] = '=' # Specifies the symbol used to separate tensor parameters from their values

Acknowledgements:

Memory functions inspired by Zach Mueller blog post. Current release_memory was adapted from the Accelerate library.