torchtrail 0.0.8

A library for tracing the execution of Pytorch operations and modules

torchtrail

torchtrail provides an external API to trace pytorch models and extract the graph of torch functions and modules that were executed. The graphs can then be visualized or used for other purposes.

Installation Instructions

On MacOs

brew install graphviz
pip install torchtrail

On Ubuntu

sudo apt-get install graphviz
pip install torchtrail

Examples

Tracing a function

import torch
import torchtrail

with torchtrail.trace():
    tensor = torch.rand(1, 64)
    tensor = torch.exp(tensor)
torchtrail.visualize(tensor, file_name="exp.svg")

Tracing a module

import torch
import transformers

import torchtrail

model_name = "google/bert_uncased_L-4_H-256_A-4"
config = transformers.BertConfig.from_pretrained(model_name)
config.num_hidden_layers = 1
model = transformers.BertModel.from_pretrained(model_name, config=config).eval()

with torchtrail.trace():
    input_tensor = torch.randint(0, model.config.vocab_size, (1, 64))
    output = model(input_tensor).last_hidden_state

torchtrail.visualize(output, max_depth=1, file_name="bert_max_depth_1.svg")

torchtrail.visualize(output, max_depth=2, file_name="bert_max_depth_2.svg")

The graph of the full module can be visualized by omitting max_depth argument

torchtrail.visualize(output, file_name="bert.svg")

Alternatively, visualization of the modules can be turned off completely using show_modules=False

torchtrail.visualize(output, show_modules=False, file_name="bert_show_modules_False.svg")