| |
|
|
| function MnistRNN() { |
| var model = this; |
|
|
| this.weights_meta = { |
| '(MnistNet).dropout(Dropout).keygen(Generator)._key': [[1973249, 1973251], [2]], |
| '(MnistNet).lstm_core(LSTMCore).fc(Linear).b': [[266496, 268544], [2048]], |
| '(MnistNet).lstm_core(LSTMCore).fc(Linear).w': [[268544, 1841408], [768, 2048]], |
| '(MnistNet).output_head(Linear).b': [[1841408, 1841665], [257]], |
| '(MnistNet).output_head(Linear).w': [[1841665, 1973249], [512, 257]], |
| '(MnistNet).pos_embed(Embed).embeddings': [[0, 200704], [784, 256]], |
| '(MnistNet).value_embed(Embed).embeddings': [[200704, 266496], [257, 256]] |
| }; |
|
|
| this.is_model_ready = false; |
| |
| this.embed_lookup = function(index, weights) { |
| return tf.slice(weights, [index], [1]); |
| }; |
|
|
| this.pos = 0; |
| this.state = null; |
| this.start_token = 256; |
| this.hidden_size = this.weights_meta['(MnistNet).lstm_core(LSTMCore).fc(Linear).b'][1][0] / 4; |
|
|
| this.initialize_state = function() { |
| this.pos = 0; |
| this.token = this.start_token; |
| var hidden = tf.zeros([1, this.hidden_size]); |
| var cell = tf.zeros([1, this.hidden_size]); |
| this.state = [hidden, cell]; |
| }; |
|
|
| this.lstm_core = function(inputs, state, weights) { |
| const [hidden, cell] = state; |
| const [w, b] = weights; |
| const i_and_h =tf.concat([inputs, hidden], 1); |
| const gated = tf.add(tf.matMul(i_and_h, w), b); |
| const [i, g, f, o] = tf.split(gated, 4, 1); |
| const f_ = tf.sigmoid(tf.add(f, 1.)); |
| const i_ = tf.sigmoid(i); |
| const g_ = tf.tanh(g); |
| const c = tf.add( |
| tf.mul(i_, g_), |
| tf.mul(cell, f_) |
| ); |
| const h = tf.mul( |
| tf.sigmoid(o), |
| tf.tanh(c) |
| ); |
| return [h, c]; |
| }; |
|
|
| this.step = function() { |
| const [token, h, c] = tf.tidy( function() { |
| const lstm_b = model.MODEL_WEIGHTS['(MnistNet).lstm_core(LSTMCore).fc(Linear).b']; |
| const lstm_w = model.MODEL_WEIGHTS['(MnistNet).lstm_core(LSTMCore).fc(Linear).w']; |
| const output_b = model.MODEL_WEIGHTS['(MnistNet).output_head(Linear).b']; |
| const output_w = model.MODEL_WEIGHTS['(MnistNet).output_head(Linear).w']; |
| const pos_embed = model.MODEL_WEIGHTS['(MnistNet).pos_embed(Embed).embeddings']; |
| const value_embed = model.MODEL_WEIGHTS['(MnistNet).value_embed(Embed).embeddings']; |
| const v = model.embed_lookup(model.token, value_embed); |
| const p = model.embed_lookup(model.pos, pos_embed); |
| const x = tf.add(v, p); |
| const [h, c] = model.lstm_core(x, model.state, [lstm_w, lstm_b]); |
| tf.dispose(model.state[0]); |
| tf.dispose(model.state[1]); |
| const logits = tf.add( |
| tf.matMul(h, output_w), |
| output_b |
| ); |
| const token = tf.multinomial(logits, 1).dataSync()[0]; |
| |
| return [token, h, c]; |
| }); |
|
|
| this.clean_memory(); |
| this.token = token; |
| this.state = [h, c]; |
| canvas.plot_xyc(this.pos, token); |
| this.pos = this.pos + 1; |
| }; |
|
|
| this.MODEL_WEIGHTS = {}; |
| this.clean_memory = function() { |
| tf.dispose(model.state[0]); |
| tf.dispose(model.state[1]); |
| }; |
|
|
| this.loop = function() { |
| this.step(); |
| if (this.pos >=28*28) { |
| setTimeout(function(){ |
| model.clean_memory(); |
| model.initialize_state(); |
| canvas.plot_grid(); |
| model.loop(); |
| }, 3000); |
| } else { |
| canvas.plot_xyc(this.pos, 255); |
| setTimeout(function(){model.loop();}, 0); |
| } |
| }; |
| |
| this.load_model_weights = function() { |
| var req = new XMLHttpRequest(); |
| req.open("GET", "weights.bin", true); |
| console.log('loading weights...'); |
| req.responseType = "arraybuffer"; |
| var this_ = this; |
| req.onload = function (event) { |
| var buff = req.response; |
| if (buff) { |
| var W = new Float32Array(buff); |
| for(var k in this_.weights_meta) { |
| info = this_.weights_meta[k]; |
| offset = info[0]; |
| shape = info[1]; |
| this_.MODEL_WEIGHTS[k] = tf.tensor(W.subarray(offset[0], offset[1]), shape); |
| } |
| this_.is_model_ready = true; |
| } else { |
| alert('Error while loading weights...'); |
| } |
| }; |
| req.send(null); |
| }; |
|
|
| this.load_when_ready = function() { |
| tf.ready().then( function() { |
| tf.enableProdMode(); |
| console.log('tf is ready'); |
| model.initialize_state() |
| model.load_model_weights(); |
| console.log(model.hidden_size); |
| }); |
| }; |
| } |
|
|
|
|
| function MnistCanvas() { |
| var canvas = document.getElementById("mnist-canvas"); |
| canvas.width = window.innerWidth; |
| canvas.height = window.innerHeight; |
| context=canvas.getContext('2d'); |
| context.translate(canvas.width/2,canvas.height/2); |
| var scale = Math.floor(Math.min(canvas.width, canvas.height) / (28*2) ) * 28; |
| console.log(scale); |
| context.scale(scale, scale) |
| context.imageSmoothingEnabled = false; |
|
|
| this.clear = function() { |
| context.clearRect(-1, -1, 2., 2.); |
| context.fillStyle = "rgb(0, 0, 0)"; |
| context.fillRect(-10, -10, 20, 20); |
| }; |
|
|
| this.plot_grid = function() { |
| for (var i=0; i< 28*28; i++) this.plot_xyc(i, 0); |
| }; |
|
|
| this.plot_xyc = function (pos, color) { |
| color = Math.max(20, color); |
| var step = 1. / 28; |
| var y = Math.floor(pos / 28 - 14) * step; |
| var x = (pos % 28 - 14) * step; |
| context.fillStyle = "rgb(0, " + color + ", 0)"; |
| context.fillRect(x, y, step, step); |
| context.strokeStyle = "rgb(0, 0, 0)"; |
| context.lineWidth = 0.008; |
| context.strokeRect(x, y, step, step); |
| }; |
|
|
| this.loading_animation = function() { |
| var counter = 0; |
| var this_ = this; |
| this_.plot_grid(); |
|
|
| var draw = function() { |
| if (model.is_model_ready) { |
| console.log('stopping animation.'); |
| model.loop(); |
| return; |
| } |
| if (counter >= 28*28) { |
| this_.plot_grid(); |
| counter = 0; |
| } |
| this_.plot_xyc(counter, 255); |
| if (counter < 28*28-1) { |
| this_.plot_xyc(counter+1, 255); |
| } |
| counter = counter+1; |
| window.requestAnimationFrame(draw); |
| }; |
| window.requestAnimationFrame(draw); |
| }; |
| } |
|
|
|
|
| var model = null; |
| var canvas = null; |
|
|
| window.onload = function() { |
| setTimeout(function() { |
| model = new MnistRNN(); |
| canvas = new MnistCanvas(); |
| console.log("init..."); |
| canvas.clear(); |
| canvas.loading_animation(); |
| model.load_when_ready(); |
| }, 500); |
| }; |
|
|
