Skip to content

A JavaScript and TypeScript port of PyTorch C++ library (libtorch) - Node.js N-API bindings for libtorch.

License

Notifications You must be signed in to change notification settings

cedrickchee/tch-js

Repository files navigation

tch-js

An unofficial JavaScript and TypeScript port of PyTorch C++ library (libtorch).

Install

Package publishing to NPM is still work-in-progress. In the meantime, you can install the package directly from GitHub.

$ npm i git+https://github.com/cedrickchee/tch-js.git

The package will download the pre-built binary during installation. You don't have to download PyTorch/libtorch or install tools for compiling from source.

Versions Supported

  • Node.js:
    • 10 (tested)
    • 12 (tested)
    • 14 (WIP)
  • PyTorch (CPU):
    • 1.4.X (tested)
    • 1.5.X
    • 1.6.X
    • 1.7.X (tested)

Code Examples

// This is a real example from an audio source separation model.
const { tch, load, Tensor } = require('tch-js');
const fs = require('fs');
const wav = require('node-wav');

// WAV samples of length 269973
const monoAudioChan = new Float32Array([
  1.100000023841858, 1.590000033378601,
  2.049999952316284, 0.18000000715255737
]);
// Flat tensor
let audio = tch.tensor(monoAudioChan); // tensor of size [269973]

// Reshape to 1xSample Length to match model input
audio = audio.view([1, monoAudioChan.length]); // tensor of size [1, 269973]

// Load PyTorch traced model async from file and return resulting ScripModule.
const model = await load('sound-model.pt');
// Forward tensor async and return resulting Tensor.
model.forward(audio, getResult);
const getResult = (err: Error, result: Tensor) => {
  if (err) return;

  // result is a tensor of size [1, 1, 269973]
  const out = result.toFloat32Array(); // convert Tensor to JS TypedArray

  // Encode output to 16-bit float WAV and write to file.
  const buf = wav.encode([out], { sampleRate: 44100, float: false, bitDepth: 16});
  fs.writeFileSync("out.wav", Buffer.from(buf));
};

Build It Yourself

Currently, only Linux builds are available.

If you want to build the package youself, below are the steps to reproduce the build.

Installing on Linux:

  • Ubuntu 18.04 (tested)
  • Ubuntu 20.04 (tested)
  1. Install build tools
$ apt install -y cmake make gcc-c++ unzip
  1. Install Node.js 10
  2. Download libtorch Download libtorch pre-cxx11 ABI and CPU version.
$ curl -o libtorch.zip https://download.pytorch.org/libtorch/cpu/libtorch-shared-with-deps-1.7.1%2Bcpu.zip
$ unzip libtorch.zip
  1. Install Node.js packages
$ npm i --ignore-scripts
  1. Build
$ npm run pre-build
  1. Test
$ npm run test

The Plan

The library should:

  • expose more libtorch types and APIs for inference
  • supports Windows
  • auto build binaries using CI
  • TypeScript types

Research

About

A JavaScript and TypeScript port of PyTorch C++ library (libtorch) - Node.js N-API bindings for libtorch.

Resources

License

Stars

Watchers

Forks

Packages

No packages published