01.py¶
```python
import torch¶
import torch.fx¶
class MyModule(torch.nn.Module):¶
def init(self):¶
super().init()¶
self.param = torch.nn.Parameter(torch.rand(3, 4))¶
self.linear = torch.nn.Linear(4, 5)¶
def forward(self, x):¶
return torch.topk(torch.sum(¶
self.linear(x + self.linear.weight).relu(), dim=-1), 3)¶
m = MyModule()¶
gm = torch.fx.symbolic_trace(m)¶
print(gm.graph)¶
gm.graph.print_tabular()¶
print(gm.code)¶
from torch.fx import symbolic_trace¶
from torch.fx.passes.graph_drawer import FxGraphDrawer¶
gm = symbolic_trace(m)¶
drawer = FxGraphDrawer(gm, "my_graph")¶
dot = drawer.get_dot_graph()¶
dot.write_svg("01.svg")¶
import torch import torch.fx
class MyModel(torch.nn.Module): def init(self): super().init() self.linear = torch.nn.Linear(4, 8)
def forward(self, x):
return torch.relu(self.linear(x))
gm = torch.fx.symbolic_trace(MyModel()) gm.graph.print_tabular()
print(gm.graph)
print(gm.code) # prints the actual code```