torchviz_model() — pytorch Function Reference

Architecture documentation for the torchviz_model() function in distributed.py from the pytorch codebase.

Function python DynamoBenchmarks Microbenchmarking called by 1

Entity Profile

DynamoBenchmarks→ Microbenchmarking→ torchviz_model() — pytorch Function Reference

Dependency Diagram

graph TD
  25f0a1d2_c9d8_0637_551b_ccf12d418bf9["torchviz_model()"]
  24e4ab44_f8ed_b835_f57a_8a13ebaabe74["run_model()"]
  24e4ab44_f8ed_b835_f57a_8a13ebaabe74 -->|calls| 25f0a1d2_c9d8_0637_551b_ccf12d418bf9
  style 25f0a1d2_c9d8_0637_551b_ccf12d418bf9 fill:#6366f1,stroke:#818cf8,color:#fff

Relationship Graph

Source Code

benchmarks/dynamo/distributed.py lines 24–32

def torchviz_model(args, model, inputs, rank):
    from torchviz import make_dot

    outputs = model(*inputs)
    loss = reduce_to_scalar_loss(outputs)
    parameter_names = dict(model.named_parameters())
    dot = make_dot(loss, params=parameter_names, show_attrs=True, show_saved=True)
    if rank == 0:
        dot.render("torchviz.dot")