Resurrect JIT functionality and loading of models

[dsyme]

The PyTorch C++ guide shows that loading of modules means using the JIT loading: https://pytorch.org/tutorials/advanced/cpp_export.html

The JIT functionality was removed after the extensive churn in Pytorch C++ API in v1.01 --> 1.50.
It needs to be resurrected.

The more direct Load/Save that was present on modules should probably not be resurrected directly.

Also need to add ONNX load/save support

[dsyme]

@NiklasGustafsson What's your thinking on this for the roadmap for TorchSharp? It feels like it should be there?

[NiklasGustafsson]

I think this is one area that requires very careful consideration. I'm not convinced that TorchSharp needs to follow the PyTorch recipe for externalizing models exactly.

While we have spent a lot of effort making the eager evaluation of models as close as possible to what PyTorch developers are doing, we do not have to make the externalized model experience the same. For example, there may be a role for Roslyn analyzers for the C# solution, rather than a clone of the Python solution. Likewise for F# -- there may be a F#-specific solution.

I think this should come post-v1.0, and we also need to consider how we will support ONNX export from TorchSharp. To me, they are two sides of the same coin.

[GeorgeS2019]

@dsyme @NiklasGustafsson

Currently, I could be wrong, it seems the continuity from pytorch saved model to Torchsharp is through ONNX.

There is a need for a proof of concept that the state dictionary exported (using TorshSharp exportsd.py script) from pyTorch can be loaded back to TorchSharp by having the requirement that both the TorchSharp and PyTorch models are compatible.

As Nikas has implied in a few occasions, one possible option is the TorchSharp support for ONNX import of PyTorch exported ONNX.

Once the model has been exported to ONNX, we can (theoretically) use that to recreate the model code in C# (or F#), which can then be used to load the model weights.

It seems TorchSharp ONNX import is equally and perhaps more important than TorchSharp ONNX export.

Again, I hope this issue is evaluated post-v1.0

Congratulation of rapid progress towards v1.0!

[dsyme]

@NiklasGustafsson I see for loading (though I'd still think we need to support loading of the pytorch format, if Java or Rust or Scala bindings do, we should check)

What about JIT? That's really what I was asking about, I forgot this issue referred to loading too.

[dsyme]

Here are relevant links from Rust PyTorch bindings

https://github.com/LaurentMazare/tch-rs#using-some-pre-trained-model

https://github.com/LaurentMazare/tch-rs/tree/main/examples/jit

I haven't checked how the ".ot" files are extracted. At least, we should be matching this

[dsyme]

Looks like they have load:

let model = tch::CModule::load(model_file)?;

[GeorgeS2019]

@dsyme @NiklasGustafsson

Based on new intiative from @dsyme, I recall the challenges of Netron dealing with PyTorch support.

PyTorch support is experimental. The recommended approach is to save the model to ONNX.

a .pth file might contain only tensor data or a partial module hierarchy without graph connections.

Rust binding approach

Then I look into the Rust binding approach as suggested by @dsyme

Rust seems to work only with TorchScript instead of PyTorch model

Example (tracing an existing module):

import torch
import torch.nn as nn

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv = nn.Conv2d(1, 1, 3)

    def forward(self, x):
        return self.conv(x)

n = Net()
example_weight = torch.rand(1, 1, 3, 3)
example_forward_input = torch.rand(1, 1, 3, 3)

# Trace a specific method and construct `ScriptModule` with
# a single `forward` method
module = torch.jit.trace(n.forward, example_forward_input)

# Trace a module (implicitly traces `forward`) and construct a
# `ScriptModule` with a single `forward` method
module = torch.jit.trace(n, example_forward_input)

The trained PyTorch model can then be serialized into a Torch Script file. The Torch Script file contains a description of the model architecture as well as trained weights. This file can be loaded from Rust to run inference for the saved model.

Given a "PyTorch" trained and save model, (e.g. pth or pt) the extension itself does not reveal if it is a PyTorch model or a TorchScript serialized model. (lots of discussion on that)

To address that in Rust.

• Rust controls that ONLY saved TorchScript model is used and converts that into a format (with extension "ot"), loadable by Rust.

The rust CModule is described here

However, it seems in practice, Rust goes from a save Jit model (e.g. "pt"), through an intermediate format ".nz", before ending up with the "ot" format.

This is what I could gather and I am not 100% sure if my understanding is correct.

[NiklasGustafsson]

@dsyme, @GeorgeS2019 -- we can and should look at JIT post-v1.0. In my opinion, we should decide on what is right in the context of .NET and the tooling options we have at our disposal. Tracing is simple, it's the scripting option of TorchScript that is going to take some thought and consideration.

I think ONNX is a good solution for exchange between languages (that's what it's for), but the reason for JIT support is primarily about performance and being able to deploy models for inference on a wide variety of platforms, so the JIT question is distinct from loading, IMO.

I'm working on an example (for the examples repo) with a simple ONNX -> C# converter that relies on the .NET model binary format. It should be ready in a week or so.