even for very simle wrappers around Pytorch layers the graph can get quite complex.
Is there a way to create subgraphs automatically fold those in 1 block?
example:
class Lin(nn.Module):
def __init__(self, *args, **kwargs):
super().__init__()
self.lin = nn.Linear(*args, **kwargs)
def forward(self, x):
x = x.view(x.shape[0], -1)
return self.lin(x)
model = nn.Sequential(*[ nn.Sequential(Lin(10), nn.ReLU()) for i in range(5)])
batch = torch.randn((4, 10))
print(model)
graph = hl.build_graph(model, batch)
graph.save("graph", format="pdf")
What I mean is something like:
# profile a submodule
lin_example = Lin(10, 10)
lin_graph = hl.build_graph(lin_example, batch)
# now fold all the submodules in the main graph
graph = graph.fold_submodule(lin_graph)
Hi,
even for very simle wrappers around Pytorch layers the graph can get quite complex. Is there a way to create subgraphs automatically fold those in 1 block?
example:
What I mean is something like: