def draw_graph(traced: torch.fx.GraphModule, fname: str, figname: str = "fx_graph"):
base, ext = os.path.splitext(fname)
if not ext:
ext = ".svg"
print(f"Writing FX graph to file: {base}{ext}")
g = graph_drawer.FxGraphDrawer(traced, figname)
x = g.get_main_dot_graph()
getattr(x, "write_" + ext.lstrip("."))(fname)
def draw_graph(traced: torch.fx.GraphModule,
fname: str,
figname: str = "fx_graph"):
if not fname.endswith(".svg"):
fname = fname + ".svg"
print(f"Writing FX graph to file: {fname}")
g = graph_drawer.FxGraphDrawer(traced, figname)
x = g.get_main_dot_graph()
x.write_svg(fname)
def draw_graph(traced: torch.fx.GraphModule, fname: str, figname: str = "fx_graph", clear_meta=True):
if clear_meta:
new_graph = copy.deepcopy(traced.graph)
traced = fx.GraphModule(traced, new_graph)
for node in traced.graph.nodes:
node.meta = {}
base, ext = os.path.splitext(fname)
if not ext:
ext = ".svg"
print(f"Writing FX graph to file: {base}{ext}")
g = graph_drawer.FxGraphDrawer(traced, figname)
x = g.get_main_dot_graph()
getattr(x, "write_" + ext.lstrip("."))(f"{base}{ext}")