pytorch

PyTorch-like deep learning library for JavaScript

A faithful JavaScript port of the PyTorch API — tensors, autograd, neural network layers, and optimizers — designed so Python/PyTorch developers feel immediately at home.

Installation

npm install pytorch

Quick Start

const { torch, nn, optim } = require('pytorch');

// Create tensors
const x = torch.tensor([[1, 2], [3, 4]]);
const y = torch.randn([2, 2]);

// Arithmetic
const z = x.add(y);
console.log(z.toString());

// Matrix multiply
const w = torch.eye(2);
console.log(x.matmul(w).toString());

Tensors

Creating Tensors

torch.tensor([[1, 2, 3], [4, 5, 6]])    // from nested array
torch.zeros([3, 4])                      // all zeros
torch.ones([2, 2])                       // all ones
torch.full([3, 3], 7.0)                  // fill with value
torch.rand([4, 4])                       // uniform [0, 1)
torch.randn([4, 4])                      // standard normal
torch.randint(0, 10, [3, 3])             // random integers
torch.eye(4)                             // identity matrix
torch.arange(0, 10, 2)                   // [0, 2, 4, 6, 8]
torch.linspace(0, 1, 5)                  // [0, 0.25, 0.5, 0.75, 1]
torch.zerosLike(existingTensor)
torch.onesLike(existingTensor)

Tensor Properties

const t = torch.randn([3, 4]);
t.shape   // [3, 4]
t.ndim    // 2
t.numel() // 12
t.sizeStr() // "torch.Size([3, 4])"

Reshaping

t.reshape([2, 6])
t.view(2, 6)
t.flatten()
t.t()        // transpose (2-D only)

Math Operations

// Element-wise
a.add(b)  // or a.add(scalar)
a.sub(b)
a.mul(b)
a.div(b)
a.pow(2)
a.neg()
a.abs()

// Reductions
t.sum()
t.mean()
t.max()
t.min()

// Matrix
a.matmul(b)   // also: a.mm(b) or torch.matmul(a, b)

Activation Functions

t.relu()
t.sigmoid()
t.tanh()
t.softmax()    // also torch.softmax(t)

Comparisons

a.eq(b)  // equal
a.ne(b)  // not equal
a.lt(b)  // less than
a.le(b)  // less or equal
a.gt(b)  // greater than
a.ge(b)  // greater or equal

Autograd

Enable gradient tracking with requiresGrad: true:

const x = torch.tensor([2.0, 3.0], { requiresGrad: true });
const y = x.mul(x).sum();  // y = x₀² + x₁²
y.backward();

console.log(x.grad.toArray()); // [4.0, 6.0]  (dy/dx = 2x)

Zero gradients before each step:

optimizer.zeroGrad();   // or model.zeroGrad()

Neural Networks (nn)

Layers

| Layer | Description | |---|---| | nn.Linear(in, out) | Fully-connected layer | | nn.ReLU() | ReLU activation | | nn.Sigmoid() | Sigmoid activation | | nn.Tanh() | Tanh activation | | nn.Softmax(dim) | Softmax | | nn.Dropout(p) | Dropout regularization | | nn.BatchNorm1d(n) | 1-D batch normalization | | nn.Sequential(...) | Layer container |

Building a Network

const model = new nn.Sequential(
  new nn.Linear(784, 256),
  new nn.ReLU(),
  new nn.Dropout(0.3),
  new nn.Linear(256, 128),
  new nn.ReLU(),
  new nn.Linear(128, 10),
);

console.log(model.toString());
// Sequential(
//   (0): Linear(in_features=784, out_features=256, bias=true)
//   (1): ReLU()
//   (2): Dropout(p=0.3)
//   ...
// )

Forward Pass

const x   = torch.randn([1, 784]);
const out = model.forward(x);
console.log(out.shape);  // [1, 10]

Loss Functions

const criterion = new nn.MSELoss();
const loss = criterion.forward(predictions, targets);

// Other losses
new nn.BCELoss()           // binary cross-entropy
new nn.CrossEntropyLoss()  // multi-class

Functional API

const { functional: F } = nn;

F.relu(x)
F.sigmoid(x)
F.softmax(x)
F.mseLoss(pred, target)
F.l1Loss(pred, target)
F.binaryCrossEntropy(pred, target)
F.crossEntropyLoss(logits, labels)

Optimizers (optim)

Available Optimizers

| Optimizer | Description | |---|---| | optim.SGD | Stochastic Gradient Descent with optional momentum & Nesterov | | optim.Adam | Adaptive Moment Estimation | | optim.AdamW | Adam with decoupled weight decay | | optim.RMSprop | Root Mean Square Propagation |

const optimizer = new optim.Adam(model.parameters(), { lr: 1e-3 });

// Training loop
for (let epoch = 0; epoch < 100; epoch++) {
  optimizer.zeroGrad();
  const out  = model.forward(x);
  const loss = criterion.forward(out, y);
  loss.backward();
  optimizer.step();
}

SGD Options

new optim.SGD(params, {
  lr:          0.01,
  momentum:    0.9,
  weightDecay: 1e-4,
  nesterov:    true,
})

Adam / AdamW Options

new optim.Adam(params, {
  lr:          1e-3,
  beta1:       0.9,
  beta2:       0.999,
  eps:         1e-8,
  weightDecay: 0,
})

Learning Rate Schedulers

const scheduler = new optim.StepLR(optimizer, 10, 0.1);
// Multiply lr by 0.1 every 10 epochs

const cosScheduler = new optim.CosineAnnealingLR(optimizer, 100);

// Call after each epoch
scheduler.step();

Full Training Example

const { torch, nn, optim } = require('pytorch');

// XOR dataset
const X = torch.tensor([[0,0], [0,1], [1,0], [1,1]]);
const y = torch.tensor([[0], [1], [1], [0]]);

// Model
const model = new nn.Sequential(
  new nn.Linear(2, 4),
  new nn.Tanh(),
  new nn.Linear(4, 1),
  new nn.Sigmoid(),
);

const criterion = new nn.MSELoss();
const optimizer = new optim.Adam(model.parameters(), { lr: 0.05 });

// Train
for (let epoch = 1; epoch <= 2000; epoch++) {
  optimizer.zeroGrad();
  const pred = model.forward(X);
  const loss = criterion.forward(pred, y);
  loss.backward();
  optimizer.step();

  if (epoch % 200 === 0)
    console.log(`Epoch ${epoch}  loss=${loss.item().toFixed(6)}`);
}

// Predict
const out = model.forward(X);
console.log('Predictions:', out.toArray());

API Reference

torch

| Function | Signature | |---|---| | tensor | (data, opts?) → Tensor | | zeros / ones | (shape, opts?) → Tensor | | rand / randn | (shape, opts?) → Tensor | | randint | (low, high, shape) → Tensor | | eye | (n) → Tensor | | arange | (start?, end, step?) → Tensor | | linspace | (start, end, steps) → Tensor | | stack / cat | (tensors, dim?) → Tensor | | exp / log / sqrt / abs | (t) → Tensor | | sum / mean | (t) → Tensor | | matmul / mm | (a, b) → Tensor | | relu / sigmoid / tanh / softmax | (t) → Tensor |

License

MIT © pytorch-js contributors