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.tlog(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:

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['max_values'] = 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 are inspired by Zach Mueller's blog post. The current release_memory function is adapted from the Accelerate library.