1. data

2. predict function

3. optimizer with learning rate for loss function

4. visualize it by line

​

*/

​

let x_data = [];

let y_data = [];

​

// function f(x) = ax + b

let a, b;

const learning_rate = 0.1;

const optimizer = tf.train.sgd(learning_rate);

let dragging = false;

​

function setup() {

  createCanvas(400, 400);

  // init random value for a and b variable

  a = tf.scalar(random(1)).variable();

  b = tf.scalar(random(1)).variable();

}

​

function mousePressed() {

  dragging = true;

}

​

function mouseReleased() {

  dragging = false;

}

​

function predict(x) {

  // this mean return f(x) = x*a + b

  return tf.tensor1d(x).mul(a).add(b);

}

​

function loss(pred, label) {

  return pred.sub(label).square().mean();

}

​

function draw() {

  background(0);

  stroke(255);

  strokeWeight(8);

  if (dragging) {

    // get the training data

    let xs = map(mouseX, 0, width, 0, 1);

    let ys = map(mouseY, 0, height, 1, 0);

    x_data.push(xs);

    y_data.push(ys);

  } else {

    if (x_data.length > 0) {

      // optimizer it

      optimizer.minimize(() => loss(predict(x_data), tf.tensor1d(y_data)), false, [a, b]);

    }

  }

  // draw all the points from data list

  for (let i = 0; i < x_data.length; i++) {

    let px = map(x_data[i], 0, 1, 0, width);

    let py = map(y_data[i], 0, 1, height, 0);

    point(px, py);

  }

  // visualize by the line

  let xs = [0, 1];

  let lineY = tf.tidy(() => predict(xs));

  let ys = lineY.dataSync();

  lineY.dispose();

  let px = map(xs[0], 0, 1, 0, width);

  let px1 = map(xs[1], 0, 1, 0, width);

​

  let py = map(ys[0], 0, 1, height, 0);

  let py1 = map(ys[1], 0, 1, height, 0);

  strokeWeight(4);

  line(px, py, px1, py1);

  console.log('numTensors (outside tidy): ' + tf.memory().numTensors);

}