Export Multiple Functions of a Pytorch Module

torch.compile
1

With torch.jit.script it was possible to export multiple functions of a pytorch module by annotating the relevant functions with torch.jit.export like so:

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

Exporting multiple functions avoids duplication when several functions go through similar computations (as in the example above with self.network).

Is (will?) this also possible with torch.export? I presume not because the functionality is undocumented, and export seems to trace individual functions, but I wanted to ask.

2

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

3

Thanks, that’s very much appreciated (and what I expected).

4

Hi, I work on export! +1 to what Soumith said – right now it’s not possible for one call of torch.export 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.

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 torch.export 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 here. Later when ExecuTorch serializes to a binary, the weights/buffer in that structure are then merged into one state dict here.

5

Thanks a lot, @angelayi for the detailed response. Your presentation at the developer conference about Executorch was great.

I have been meaning to get in ExecuTorch for some time, and this seems to be a great excuse.

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?

1 Like
6

no, ExecuTorch is only for embedded device runtimes, so it wont be supporting CUDA delegates

7

Thanks fo your reply, @smth .

8

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.

9

Hi @smth , @angelayi any plans of addressing this first class in torch.export?

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).

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?