Export Multiple Functions of a Pytorch Module

fabian_schutze: <p>With torch.jit.script it was possible to export multiple functions of a pytorch module by annotating the relevant functions with <code>torch.jit.export</code> like so:</p>
<pre><code class="lang-auto">import torch
class Module(torch.nn.Module):

    def __init__(self):
        super().__init__()
        self.network = torch.nn.Sequential(torch.nn.Identity())

    @torch.jit.export
    def forward(self, tensor_foo):
        return self.network(tensor_foo)

    @torch.jit.export
    def bar(self, tensor_bar):
        first_step = self.network(tensor_bar)
        second_step = torch.nn.functional.sigmoid(first_step)
        return second_step
</code></pre>
<p>Exporting multiple functions avoids duplication when several functions go through similar computations (as in the example above with <code>self.network</code>).</p>
<p>Is (will?) this also possible with <a href="https://pytorch.org/docs/stable/export.html" rel="noopener nofollow ugc">torch.export</a>? I presume not because the functionality is undocumented, and export seems to trace individual functions, but I wanted to ask.</p>

smth: <p><code>torch.export</code> traces paths, and the automatically pulls relevant source code (+ inserts guards for types). In that sense, it isn’t doing anything like <code>torch.jit.script</code> and wouldn’t have an equivalent to <code>torch.jit.export</code>. That being said, I’ll ask someone who’s more of an expert to double-check my answer.</p>

fabian_schutze: <p>Thanks, that’s very much appreciated (and what I expected).</p>

angelayi: <p>Hi, I work on export! +1 to what Soumith said – right now it’s not possible for one call of <code>torch.export</code> to trace multiple functions or avoid duplicate work when functions go through similar computation. It just goes down the given function and traces through any relevant code.</p>
<p>This is not exactly what you asked about, but for ExecuTorch we have a concept of “multiple entry-point” support in that you can call <code>torch.export</code> multiple times, and then combine them into one structure and pass it throughout ExecuTorch. The API has since gone through some refactoring, but you can find the original implementation <a href="https://github.com/pytorch/executorch/blob/5bc506624bf0d08a0791c142215795dc812e2120/exir/capture/_capture.py#L321" rel="noopener nofollow ugc">here</a>. Later when ExecuTorch serializes to a binary, the weights/buffer in that structure are then merged into one state dict <a href="https://github.com/pytorch/executorch/blob/5bc506624bf0d08a0791c142215795dc812e2120/exir/emit/_emitter.py#L371-L378" rel="noopener nofollow ugc">here</a>.</p>

fabian_schutze: <p>Thanks a lot, <a class="mention" href="/u/angelayi">@angelayi</a>  for the detailed response. Your presentation at the developer conference about Executorch was great.</p>
<p>I have been meaning to get in ExecuTorch for some time, and this seems to be a great excuse.</p>
<p>I didn’t fully get whether Executorch will support a CUDA delegate. I didn’t see that on the slides from last year’s developer conference and as TF lite doesn’t support it, I’m wondering whether Executorch will support it?</p>

smth: <p>no, ExecuTorch is only for embedded device runtimes, so it wont be supporting CUDA delegates</p>

fabian_schutze: <p>Thanks fo your reply, <a class="mention" href="/u/smth">@smth</a> .</p>

angelayi: <p>I think theoretically it could support CUDA delegates, although no work has been done on it yet because the main purpose is for embedded device runtimes.</p>

raziel: <p>Hi <a class="mention" href="/u/smth">@smth</a> , <a class="mention" href="/u/angelayi">@angelayi</a>  any plans of addressing this first class in torch.export?</p>
<p>It’s very useful to have multiple entry points bundled within the exported program (hence why we did it in Executorch). The executorch implementation was a bit of a “we need it soon, so let’s do it at this level” but it would be cleaner to have this being supported from the top (i.e. torch.export).</p>
<p>We had iterated over a number of solutions, and the executorch one seemed like a nice tradeoff (certainly at least a good top level API, even if the implementation may be different to allow for easier debugging across the multiple entry points export issues). Maybe that API can be promoted?</p>