repository·main·Indexed 22 days ago
https://github.com/eduardoleao052/js-pytorchA Deep Learning JavaScript library built from scratch to closely follow PyTorch's syntax. It supports neural network construction and training in Node.js and browser environments, featuring GPU support via GPU.js, automatic differentiation (autograd), and a comprehensive set of tools including the torch.nn namespace for layers (Linear, Transformer Block, Embedding), torch.optim for optimizers like Adam, and utilities for saving and loading models via JSON.
To install the library locally for use in Node.js environments on MacOS, Windows, or Ubuntu, use the following command:
npm install js-pytorchNote for Windows users: If you encounter errors during installation, you may need to install the latest version of Visual Studio with the "Desktop development with C++" workload included.
To use the library in a web browser, include the following script tag in the <head> of your HTML file. You can then use the torch global object within any <script> tag in your HTML body.
<head>
<title>My Project</title>
<script src="https://cdnjs.cloudflare.com/ajax/libs/js-pytorch/0.7.2/js-pytorch-browser.js"
integrity="sha512-l22t7GnqXvHBMCBvPUBdFO2TEYxnb1ziCGcDQcpTB2un16IPA4FE5SIZ8bUR+RwoDZGikQkWisO+fhnakXt9rg=="
crossorigin="anonymous"
referrerpolicy="no-referrer">
</script>
</head>
<body>
<script>
let x = torch.randn([10,5])
let linear = new torch.nn.Linear(5,1,'gpu',true)
let z = linear.forward(x)
console.log(z.data)
</script>
</body><head>
<title>My Project</title>
<!-- New script goes here -->
<script src="https://cdnjs.cloudflare.com/ajax/libs/js-pytorch/0.7.2/js-pytorch-browser.js"
integrity="sha512-l22t7GnqXvHBMCBvPUBdFO2TEYxnb1ziCGcDQcpTB2un16IPA4FE5SIZ8bUR+RwoDZGikQkWisO+fhnakXt9rg=="
crossorigin="anonymous"
referrerpolicy="no-referrer">
</script>
<!---->
</head>
<body
<script>
let x = torch.randn([10,5])
let linear = new torch.nn.Linear(5,1,'gpu',true)
let z = linear.forward(x)
console.log(z.data)
</script>
</body>The Module class is the base class for all neural network layers and models. It manages learnable parameters and tracks the operational mode (train or eval).
Key capabilities:
.parameters() to retrieve a flat list of all Parameter and Tensor objects within the module (and its sub-modules) that require gradients..train() to set the module and all its sub-modules to training mode, or .eval() to set them to evaluation mode. This is critical for layers like Dropout that behave differently during training vs. inference..entries() to get an array of [key, value] pairs for all enumerable properties of the module.To build complex models, extend nn.Module and implement the forward(x) method. Training involves a loop of:
nn.CrossEntropyLoss).loss.backward() to backpropagate.optimizer.step() to update weights.optimizer.zero_grad() to reset gradients.// Require the Library if running in node (not necessary in the browser):
const { torch } = require("js-pytorch");
const nn = torch.nn;
const optim = torch.optim;
const device = 'gpu';
// Define training hyperparameters:
const vocab_size = 52;
const hidden_size = 32;
const n_timesteps = 16;
const n_heads = 4;
const dropout_p = 0;
const batch_size = 8;
// Create Transformer decoder Module:
class Transformer extends nn.Module {
constructor(vocab_size, hidden_size, n_timesteps, n_heads, dropout_p, device) {
super();
// Instantiate Transformer's Layers:
this.embed = new nn.Embedding(vocab_size, hidden_size);
this.pos_embed = new nn.PositionalEmbedding(n_timesteps, hidden_size);
this.b1 = new nn.Block(hidden_size, hidden_size, n_heads, n_timesteps, dropout_p, device);
this.b2 = new nn.Block(hidden_size, hidden_size, n_heads, n_timesteps, dropout_p, device);
this.ln = new nn.LayerNorm(hidden_size);
this.linear = new nn.Linear(hidden_size, vocab_size, device);
}
forward(x) {
let z;
z = torch.add(this.embed.forward(x), this.pos_embed.forward(x));
z = this.b1.forward(z);
z = this.b2.forward(z);
z = this.ln.forward(z);
z = this.linear.forward(z);
return z;
}
}
// Instantiate your custom nn.Module:
const model = new Transformer(vocab_size, hidden_size, n_timesteps, n_heads, dropout_p, device);
// Define loss function and optimizer:
const loss_func = new nn.CrossEntropyLoss();
const optimizer = new optim.Adam(model.parameters(), (lr = 5e-3), (reg = 0));
// Instantiate sample input and output:
let x = torch.randint(0, vocab_size, [batch_size, n_timesteps, 1]);
let y = torch.randint(0, vocab_size, [batch_size, n_timesteps]);
let loss;
// Training Loop:
for (let i = 0; i < 40; i++) {
// Forward pass through the Transformer:
let z = model.forward(x);
// Get loss:
loss = loss_func.forward(z, y);
// Backpropagate the loss using torch.tensor's backward() method:
loss.backward();
// Update the weights:
optimizer.step();
// Reset the gradients to zero after each training step:
optimizer.zero_grad();
// Print loss at every iteration:
console.log(`Iter ${i} - Loss ${loss.data[0].toFixed(4)}`)
}You can persist models to a JSON file using torch.save and reload them using torch.load. When loading, you must first instantiate a placeholder object using the same architecture as the original model to receive the loaded weights.
// Instantiate your model:
const model = new Transformer(vocab_size, hidden_size, n_timesteps, n_heads, dropout_p);
// Train the model:
trainModel(model);
// Save model to JSON file:
torch.save(model, 'model.json')
// To load, instantiate placeHolder using the original model's architecture:
const placeHolder = new Transformer(vocab_size, hidden_size, n_timesteps, n_heads, dropout_p);
// Load weights into placeHolder:
const newModel = torch.load(placeHolder, 'model.json')You can perform automatic differentiation (autograd) by defining operations on Tensors and calling the .backward() method on the result. You can specify a device (e.g., 'gpu') when instantiating Tensors or nn.Modules. To track gradients, pass true as the second argument to the Tensor constructor.
// Require the Library if running in node (not necessary in the browser):
const { torch } = require("js-pytorch");
// Pass device as an argument to a Tensor or nn.Module (same as PyTorch):
const device = 'gpu';
// Instantiate Tensors:
let x = torch.randn([8, 4, 5]);
let w = torch.randn([8, 5, 4], true, device);
let b = torch.tensor([0.2, 0.5, 0.1, 0.0], true);
// Make calculations:
let out = torch.matmul(x, w);
out = torch.add(out, b);
// Compute gradients on whole graph:
out.backward();
// Get gradients from specific Tensors:
console.log(w.grad);
console.log(b.grad);The following commands are available for developers working on the library:
npm run build (outputs CJS, ESM, and index.d.ts to dist/)npm run lint (uses ESLint)npm testnpm run prettier (uses Prettier)npm run benchnpm run bench:updateBenchmarks are located in the tests/benchmarks/ directory.
Create a new Tensor instance by providing an iterable (like an array) or a single number. You can optionally specify if the tensor should track gradients for autograd and which device to use.
Parameters:
data: Array<any> | number - The data to be stored.requires_grad: boolean (default: false) - Whether to track gradients for this tensor.device: string (default: 'cpu') - The device to store the tensor on ('cpu' or 'gpu').Advanced indexing and conditional filling:
at(index1, index2?)
Extracts elements using index lists. index1 is required; index2 is optional for 2D extraction.
masked_fill(mask, condition, value)
Fills elements in the tensor with value where the condition function returns true for the corresponding element in the mask tensor.
Broadcasting allows performing operations between tensors of different shapes by expanding or contracting dimensions.
broadcast(a, b): Broadcasts tensor a into the shape of tensor b. If the shape of a is smaller, it may be expanded; if larger, it may be summed.broadcastUp(inElement, outElement): Adds new dimensions to inElement until its depth matches outElement.The library provides several functional interfaces for common tensor operations. Most of these can be called as standalone functions passing a Tensor as the first argument:
log(a): Element-wise natural logarithm.matmul(a, b): Matrix multiplication of tensors a and b.transpose(a, dim1, dim2): Transposes the tensor along two consecutive dimensions.reshape(a, shape): Reshapes the tensor into the specified shape array. The total number of elements must remain constant.at(a, idx1, [idx2]): Extracts elements from the tensor using indices. idx1 and idx2 can be Tensor or Array.masked_fill(a, mask, condition, value): Fills elements in a where the mask satisfies the condition function with the specified value.torch object is the primary entrypoint for the library. It provides access to tensor operations, neural network layers (nn), optimizers (optim), and utility functions. Most operations are accessed via the torch namespace.