TorchJS is a JS binding for PyTorch. Its primary objective is to allow running Torch Script inside Node.js program. Complete binding of libtorch is possible but is out-of-scope at the moment.
-
Add support for
List(JavascriptArray),Dict(JavascriptObject),String,float(Javascriptnumber) as inputs and outputs. -
Add CUDA support.
-
Add ops from torchvision.
-
Add async support for
forwardfunction.
To install the forked version, you can install it from npm:
yarn add torch-js@npm:@arition/torch-jsIn tests/resources/torch_module.py, you will find the defination of our test module and the code to generate the trace file.
class TestModule(torch.nn.Module):
def __init__(self):
super(TestModule, self).__init__()
def forward(self, input1, input2):
return input1 + input2Once you have the trace file, it may be loaded into NodeJS like this
const torch = require("torch-js");
const modelPath = `test_model.pt`;
const model = new torch.ScriptModule(testModelPath);
const inputA = torch.rand([1, 5]);
const inputB = torch.rand([1, 5]);
const res = await model.forward(inputA, inputB);More examples regarding tensor creation, ScriptModule operations, and loading models can be found in our examples folder.