Skip to content

Commit

Permalink
Merge pull request #26 from eduardoleao052/gpu
Browse files Browse the repository at this point in the history 
Gpu support
  • Loading branch information
@eduardoleao052
eduardoleao052 authored Jul 19, 2024
2 parents 8af32df + a516d99 commit 3d06455
Show file tree
Hide file tree
Showing 12 changed files with 505 additions and 159 deletions.
35 changes: 28 additions & 7 deletions README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,8 @@
# PyTorch in JavaScript

- JS-PyTorch is a Deep Learning **JavaScript library** built from scratch, to closely follow PyTorch's syntax.
- It contains a fully functional [Tensor](src/tensor.ts) object, which can track gradients, Deep Learning [Layers](src/layers.ts) and functions, and an **Automatic Differentiation** engine.
- This library has **GPU support**, using GPU.js.
- It contains a gradient-tracking [Tensor](src/tensor.ts) object, Deep Learning [Layers](src/layers.ts) and functions, and an **Automatic Differentiation** engine.
- Feel free to try out the <a href="https://eduardoleao052.github.io/js-pytorch/assets/demo/demo.html" target="blank">Web Demo</a>!

> **Note:** You can install the package locally with: `npm install js-pytorch`
Expand Down Expand Up @@ -94,7 +95,7 @@ const { torch } = require("js-pytorch");

// Instantiate Tensors:
let x = torch.randn([8, 4, 5]);
let w = torch.randn([8, 5, 4], (requires_grad = true));
let w = torch.randn([8, 5, 4], (requires_grad = true), (device = 'gpu));
let b = torch.tensor([0.2, 0.5, 0.1, 0.0], (requires_grad = true));

// Make calculations:
Expand All @@ -115,6 +116,7 @@ console.log(b.grad);
const { torch } = require("js-pytorch");
const nn = torch.nn;
const optim = torch.optim;
const device = 'gpu';

// Define training hyperparameters:
const vocab_size = 52;
Expand All @@ -126,15 +128,15 @@ const batch_size = 8;

// Create Transformer decoder Module:
class Transformer extends nn.Module {
constructor(vocab_size, hidden_size, n_timesteps, n_heads, dropout_p) {
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);
this.b2 = new nn.Block(hidden_size, hidden_size, n_heads, n_timesteps,dropout_p);
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);
this.linear = new nn.Linear(hidden_size, vocab_size, device);
}

forward(x) {
Expand All @@ -149,7 +151,7 @@ class Transformer extends nn.Module {
}

// Instantiate your custom nn.Module:
const model = new Transformer(vocab_size, hidden_size, n_timesteps, n_heads, dropout_p);
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();
Expand Down Expand Up @@ -182,6 +184,25 @@ for (let i = 0; i < 40; i++) {
}
```
### Saving and Loading models:
```typescript
// 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')
```
<br/>
## 3. Distribution & Devtools
Expand Down
Binary file added assets/demo/bundle.js
View file
Binary file not shown.
15 changes: 14 additions & 1 deletion assets/demo/demo.css
View file
Original file line number Diff line number Diff line change
Expand Up @@ -141,6 +141,20 @@ button:active {
background-color: #821414;
}

.device-button {
width: 10%;
background-color: #818181;
padding-left: 1%;
}

.device-button:hover {
background-color: #5f5f5f;
}

.device-button:active {
background-color: #3a3a3a;
}

.icon {
width: 72%;
margin: auto;
Expand Down Expand Up @@ -201,4 +215,3 @@ input {
font-size: 16px;
color:#1b1b1b;
}

87 changes: 70 additions & 17 deletions assets/demo/demo.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,28 +6,28 @@
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>Neural Network Demo</title>

<script type="module" src="index.mjs"></script>
<script src="bundle.js"></script>
<script type="module" src="index.js"></script>

<link rel="script" href="demo.js">
<link rel="stylesheet" href="demo.css">
</head>

<body>
<script type="module">
import { torch } from './index.mjs';
import { torch } from './index.js';
window.torch = torch;
</script>

<div style="height: 20px;"></div>
<div class="container">
<div class="separator" id="right-separator" style="display: block; padding-bottom: 0px;">
<img style="display: block;" class='icon' src="../../assets/icon_gray_bg.png">
<img style="display: block;" class='icon' src="../../assets/demo_logo.png">
<a target="_blank" href="https://github.com/eduardoleao052/js-pytorch"> <img class="logos" src="../../assets/github.png" style="width: 25px; height: 25px; display: inline-block; margin-bottom: 5px;"></a>
<a target="_blank" href="https://www.linkedin.com/in/eduardo-leao-368bb6259/"> <img class="logos" src="../../assets/linkedin.png" style="width: 25px; height: 25px; display: inline-block; margin-bottom: 5px;"></a>
</div>

<div class="separator" id="right-separator" style="margin-top: 10px; display: block;">
<p style="display: block; font-size: 16px; color: #292929; margin-bottom: 25px;">Welcome to <b>JS-Torch's Web Demo</b>! You can choose the <b>Model Hyperparameters</b> on the left, set the <b>Model Layers</b> on the right (number of layers and hidden dimension on each). </p>
<p style="display: block; font-size: 16px; color: #292929; margin-bottom: 25px;">Welcome to <b>JS-PyTorch's Web Demo</b>! You can choose the <b>Model Hyperparameters</b> on the left, set the <b>Model Layers</b> on the right (number of layers and hidden dimension on each). </p>
<ul>
<li>This model was trained on a <b>Dummy dataset</b>, composed of <b>randomly generated images</b>. </li>
<br>
Expand Down Expand Up @@ -78,6 +78,9 @@ <h2 style="display: inline-block; margin-right: 15px; margin-left: 20px;">Model
<button id="start-button" onclick="trainLoopInitializer()">Train</button>
<button id="stop-button" onclick="pauseTraining()">Pause</button>
<button id="reset-button" onclick="resetTraining()">Reset</button>
<div style="width: 25px; display: inline-block;"></div>
<button id="cpu-trigger" class="device-button" onclick="CPUCaller()">CPU</button>
<button id="gpu-trigger" class="device-button" onclick="GPUCaller()">GPU</button>
</div>
</div>
</div>
Expand All @@ -90,13 +93,14 @@ <h2>Graph</h2>
<p id="iter"> <b>Iteration:</b> </p>
<p id="total-visited"> <b>Total Training Examples:</b> </p>
<p id="loss"> <b>Loss:</b> </p>
<p id="device-showcase"> <b>Device:</b> </p>
</div>
<divs tyle="display: inline-block;">
<a target="_blank" href="https://github.com/eduardoleao052/js-pytorch"><img class="logos" src="../../assets/github.png" style="width: 25px; height: 25px; display: inline-block; margin-bottom: 0px;"></a>
<a target="_blank" href="https://www.linkedin.com/in/eduardo-leao-368bb6259/"><img class="logos" src="../../assets/linkedin.png" style="width: 25px; height: 25px; display: inline-block; margin-bottom: 0px;"></a>
</div>
</div>

<script src="./gpu-browser.min.js"></script>
<script>
let boxCount = [];
let data = [];
Expand All @@ -105,6 +109,43 @@ <h2>Graph</h2>
let overFlow = 1;
let iter = 0;
let total_visited = 0;
let device = 'cpu';

function doMytest() {
const gpu = new GPU.GPU();
const kernelFunc = function(a, b) {
let sum2 = 0;
for (let i = 0; i < len; i++) {
sum2 += a[this.thread.y][i] * b[i][this.thread.x];
}
return sum2;
};
let forwardKernel = gpu.createKernel(kernelFunc, { loopMaxIterations: 512 }).setOutput([512, 512]);
let a = torch.randn([512,512])
let b = torch.randn([512,512])
let c = forwardKernel(a.data, b.data)
console.log(c)
}

function GPUCaller() {
if (!training) {
device = 'gpu';
let _cpu = document.getElementById('cpu-trigger');
let _gpu = document.getElementById('gpu-trigger');
_gpu.style.backgroundColor = '#3e3e3e';
_cpu.style.backgroundColor = '';
}
}
function CPUCaller() {
if (!training) {
device = 'cpu';
let _cpu = document.getElementById('cpu-trigger');
let _gpu = document.getElementById('gpu-trigger');
_cpu.style.backgroundColor = '#3e3e3e';
_gpu.style.backgroundColor = '';
}
}


function addBox() {
if (boxCount.length < 5 && !training) {
Expand All @@ -119,7 +160,7 @@ <h2>Graph</h2>
inputField.type = 'number';
// The next Box has as many neurons as last current Box:
inputField.value = boxCount[boxCount.length -1] || '10';
inputField.max = '100';
inputField.max = '256';
inputField.idx = boxCount.length;
inputField.id = `box-input`;
inputField.onchange = function() {
Expand All @@ -140,8 +181,8 @@ <h2>Graph</h2>
};

function changeDim(el, boxCount) {
if (el.value > 100) {
el.value = 100;
if (el.value > 256) {
el.value = 256;
};
boxCount[el.idx] = Number(el.value);
};
Expand Down Expand Up @@ -183,27 +224,34 @@ <h2>Graph</h2>
for (button of buttons) {
button.style.backgroundColor = '#0056b3';
};
let _cpu = document.getElementById('cpu-trigger');
let _gpu = document.getElementById('gpu-trigger');
if (device === 'cpu') {
_cpu.style.backgroundColor = '#3e3e3e';
} else {
_gpu.style.backgroundColor = '#3e3e3e';
}
// Implement dummy torch.nn.Module class:
class NeuralNet extends torch.nn.Module {
constructor() {
super();
// Instantiate Neural Network's Layers:
this.wIn = new torch.nn.Linear(32,boxCount[0]);
this.wIn = new torch.nn.Linear(32,boxCount[0],device);
this.reluIn = new torch.nn.ReLU();
if (boxCount.length > 1) {
this.w1 = new torch.nn.Linear(boxCount[0],boxCount[1]);
this.w1 = new torch.nn.Linear(boxCount[0],boxCount[1],device);
this.relu1 = new torch.nn.ReLU();
} if (boxCount.length > 2) {
this.w2 = new torch.nn.Linear(boxCount[1],boxCount[2]);
this.w2 = new torch.nn.Linear(boxCount[1],boxCount[2],device);
this.relu2 = new torch.nn.ReLU();
} if (boxCount.length > 3) {
this.w3 = new torch.nn.Linear(boxCount[2],boxCount[3]);
this.w3 = new torch.nn.Linear(boxCount[2],boxCount[3],device);
this.relu3 = new torch.nn.ReLU();
} if (boxCount.length > 4) {
this.w4 = new torch.nn.Linear(boxCount[3],boxCount[4]);
this.w4 = new torch.nn.Linear(boxCount[3],boxCount[4],device);
this.relu4 = new torch.nn.ReLU();
};
this.wOut = new torch.nn.Linear(boxCount[boxCount.length-1], 32);
this.wOut = new torch.nn.Linear(boxCount[boxCount.length-1], 32, device);

this.reluOut = new torch.nn.ReLU();

Expand Down Expand Up @@ -285,6 +333,7 @@ <h2>Graph</h2>
document.getElementById('iter').innerHTML = `<b>Iteration:</b> ${iter}`;
document.getElementById('total-visited').innerHTML = `<b>Total Training Examples:</b> ${total_visited}`;
document.getElementById('loss').innerHTML = `<b>Loss:</b> ${loss.data[0].toFixed(3)}`;
document.getElementById('device-showcase').innerHTML = `<b>Device:</b> ${device}`;
iter += 1;
total_visited += batch_size;

Expand All @@ -307,7 +356,11 @@ <h2>Graph</h2>
for (button of buttons) {
button.style.backgroundColor = '';

};
}
buttons = document.getElementsByClassName('device-button');
for (button of buttons) {
button.style.backgroundColor = '';
}
};

function plotGraph() {
Expand Down Expand Up @@ -401,4 +454,4 @@ <h2>Graph</h2>
</script>
</body>

</html>
</html>
28 changes: 28 additions & 0 deletions assets/demo/gpu-browser.min.js
View file

Large diffs are not rendered by default.

Loading

0 comments on commit 3d06455

Please sign in to comment.