Export Multiple Functions of a Pytorch Module

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.

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.

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

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.

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

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

Thanks fo your reply, @smth .

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.