frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

add use_strict_trace to tensorboard add_graph method (#63120)

Browse Source 
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/63120

FAIM returns dictionaries as the model output, which throws an error when trying to trace using add_graph. Pass in `strict` to the tracer to make this user configurable.

User post: https://fb.workplace.com/groups/pytorchLightning/permalink/1510194972650369/?comment_id=1510252919311241&reply_comment_id=1510281112641755

Test Plan: unit test

Reviewed By: Reubend

Differential Revision: D30265890

fbshipit-source-id: 58b25d9500b875a29a664aa9ef4c1e7f13631fa1
This commit is contained in:
Edward Wang (EcoF)
2021-08-12 12:10:50 -07:00
committed by Facebook GitHub Bot
parent 1022443168
commit 96fb1a56ea
3 changed files with 45 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
torch/utils/tensorboard/writer.py
View File

@ -718,9 +718,7 @@ class SummaryWriter(object):
torch._C._log_api_usage_once("tensorboard.logging.add_onnx_graph")
self._get_file_writer().add_onnx_graph(load_onnx_graph(prototxt))
def add_graph(self, model, input_to_model=None, verbose=False):
# prohibit second call?
# no, let tensorboard handle it and show its warning message.
def add_graph(self, model, input_to_model=None, verbose=False, use_strict_trace=True):
"""Add graph data to summary.
Args:
@ -728,11 +726,14 @@ class SummaryWriter(object):
input_to_model (torch.Tensor or list of torch.Tensor): A variable or a tuple of
variables to be fed.
verbose (bool): Whether to print graph structure in console.
use_strict_trace (bool): Whether to pass keyword argument `strict` to
`torch.jit.trace`. Pass False when you want the tracer to
record your mutable container types (list, dict)
"""
torch._C._log_api_usage_once("tensorboard.logging.add_graph")
if hasattr(model, 'forward'):
# A valid PyTorch model should have a 'forward' method
self._get_file_writer().add_graph(graph(model, input_to_model, verbose))
self._get_file_writer().add_graph(graph(model, input_to_model, verbose, use_strict_trace))
else:
# Caffe2 models do not have the 'forward' method
from caffe2.proto import caffe2_pb2