Loading .pt model weights

[StudyingLover]

Describe the Feature

I have trained a PyTorch model, and I want to directly load this model and run inference in a browser. However, the current method for loading the weights requires defining the model first.

Motivation [optional]

No response

Sample Code

for example

const { torch } = require("js-pytorch");

const model = torch.load("model.pth")

Would you like to work on this issue?

None

[eduardoleao052]

I have been busy with other issues recently, but will give it a shot now. Direct compatibility with PyTorch models is not something I had originally thought of, but it's worth a try!

[ocramz]

Thank you, came here to ask the same thing.

[bitanath]

Motivated to work on this! Here's my preliminary very basic analysis and approach:

I'd like to propose a method with signature load_state_dict(state_dict:Record<string,Tensor>) added in layers.ts order to achieve this. Basically export PT weights to JSON through a minimalist python package as illustrated below:

weights_json = OrderedDict({key: state_dict[key].detach().cpu().tolist() for key in state_dict})
with open("weights.json","w") as f:
    json.dump(weights_json,f)

Following which a more complex TS parser maps the JSON layer weights and biases to the existing modules in order to recreate the js-pytorch nn module with transpositions of the weights if necessary. Since js-pytorch only supports a small subset of torch layers this can be achieved through a manual comparison of weight and bias shapes that PyTorch saves as.

I have not done this one to one comparison, but there seem to be only about 5 basic trainable layers that are analogous with PyTorch and weighted as well, namely:

• nn.Linear
• nn.MultiHeadSelfAttention
• nn.Embedding
• nn.PositionalEmbedding
• nn.LayerNorm

I do not believe other layers need to be mapped since they are non parametric learnable. However, please correct me if I'm wrong. Any guidance regarding the specifications of the typescript classes is also appreciated.

The most important use case/motivation of this is to train using CUDA accelerated PyTorch in a python notebook and to inference in browser with whatever means available

This library wins over ONNX runtime envs / torchscript wherever limited server or browser control exists, but deployment within a monorepo or monolithic environment, for eg. React based web apps on pre-built platforms or integration of a thin neural network within existing server framework.