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
number) as inputs and outputs.
Add CUDA support.
Add ops from torchvision.
Add async support for
To install the forked version, you can install it from npm:
yarn add torch-js@npm:@arition/torch-js
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 + input2
Once 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.