NN graph view using torch script

[soumyarooproy]

Referring to this:
https://pytorch.org/docs/master/jit.html#creating-torch-script-code

“It [ScriptModule] is an analogue of torch’s nn.Module and represents an entire model as a tree of submodules.”

How does a ScriptModule object get serialized using the ScriptModule.save() method? Is it still an NN graph but with special conditional nodes? Or is it more of a code AST and the graph-level abstraction is lost?

Based on the following question on the PyTorch discussion forum, tracing seems to just export to ONNX format:

How to use torch.jit for graph trace?

Hi all, I try to trace myself graph using torch.jit but i have next problem: Create simple model class SimpleNet(nn.Module): def __init__(self): super().__init__() self.conv1 = nn.Conv2d(3, 10, kernel_size=1) def forward(self, x): x = self.conv1(x) return x Create variable and trace model: net = SimpleNet() var = Variable(torch.randn(1, 3, 224, 224)) trace, out = jit.trace(net, var) trace.export(list(net.state_dict().values())) And I get next o…

[smth]

@soumyarooproy

Is it still an NN graph but with special conditional nodes? Or is it more of a code AST and the graph-level abstraction is lost?

I dont think they’re that different. Sure, the code AST means you have some constants, some scalars, some booleans, but for the large part you have Tensor computations. For example:

if dilation > 0:
    aten::dilated_convolution(...)
else:
    aten::convolution(...)

Currently torchscript has fundamental types of Tensors, tuples, lists, scalars, constants, None, and constant strings.
So you can theoretically make a python program out of it that doesn’t do neural network workloads, but for the large part you prob. will have a graph of Tensor operations.