TensorFlow.js (一)

发表于 更新于 分类于

项目地址

谷歌的TensorFlow.js示例代码中的polynomial-regression-core拟合曲线
https://github.com/tensorflow/tfjs-examples/tree/master/polynomial-regression-core

代码理解

主要代码在index.js中, 代码做的事情是随机生成一条y = a x 3 + b x 2 + c x + d的曲线A, 然后在这条曲线上下随机生成一些点做为数据集。
通过这些数据集我们要找到合适a, b, c, d。正常情况下我们找到的a, b, c, d生成的曲线和曲线A会基本重合。
数据集生成在data.js中, 以下是对index.js代码的理解

1
2
3
import * as tf from '@tensorflow/tfjs';
import { generateData } from './data';
import { plotData, plotDataAndPredictions, renderCoefficients } from './ui';

生成随机的a, b, c, d作为初始值

1
2
3
4
const a = tf.variable(tf.scalar(Math.random()));
const b = tf.variable(tf.scalar(Math.random()));
const c = tf.variable(tf.scalar(Math.random()));
const d = tf.variable(tf.scalar(Math.random()));

迭代次数, 可以自行调整, 调高可以提供准确性, 花费时间更长

1
const numIterations = 300;

学习率, 相当于每次迭代调整的步伐, 步伐小,则需要迭代更多轮, 但并不是越大越好, 太大可能一步跨过最佳的值

1
const learningRate = 0.5;

SGD优化器在下面迭代的时候用到

1
const optimizer = tf.train.sgd(learningRate);

构建模型, tf.tidy可以在方法执行之后自动清理其中的张量, 释放内存, 但是不会清除函数的返回值
这里构建了一个y = a * x ^ 3 + b * x ^ 2 + c * x + d的模型, 这里的x指的不是单个x, 而是x的矩阵, 进行矩阵运算, 得到y的矩阵

1
2
3
4
5
6
7
8
function predict(x) {
  return tf.tidy(() => {
    return a.mul(x.pow(tf.scalar(3, 'int32')))
      .add(b.mul(x.square()))
      .add(c.mul(x))
      .add(d);
  });
}

定义损失函数, 返回的值是结果与预期值的方差

1
2
3
4
function loss(prediction, labels) {
  const error = prediction.sub(labels).square().mean();
  return error;
}

循环训练模型, numIterations是开头定义的迭代次数, optimizer是SGD优化器, 这里asyncz指明train是异步函数, 每次迭代通过await tf.nextFrame();等待一次计算的结束再进行下一次计算, 对于a, b, c, d值的修改更新, 都由优化器负责

1
2
3
4
5
6
7
8
9
async function train(xs, ys, numIterations) {
  for (let iter = 0; iter < numIterations; iter++) {
    optimizer.minimize(() => {
      const pred = predict(xs);
      return loss(pred, ys);
    });
    await tf.nextFrame();
  }
}

之后的代码是生成训练数据, 运行项目, 我们可以看到三幅图, 第二副图的曲线是初始状态随机的a, b, c, d所生成的曲线, 第三幅图是训练之后的, 结果与曲线A相差无几

1
2
$ npm install
$ npm run watch