Published
Edited
Jan 12, 2019
1 fork
39 stars
tf = require('@tensorflow/tfjs')
inputData = {
// tf.tidy will clean up all the GPU memory used by tensors
// inside this function, other than the tensor that is returned.
const data = tf.tidy(() => {
const NUM_POINTS = 200;
const [a, b, c, d] = coeff.map(i => tf.scalar(i)); // a = tf.scalar(coeff[0]), and so on
const x = tf.randomUniform([NUM_POINTS], -1, 1);
const y = a.mul(x.pow(tf.scalar(3)))
.add(b.mul(x.square()))
.add(c.mul(x))
.add(d)
.add(tf.randomNormal([NUM_POINTS], 0, sd));
return {x, y};
});
yield data;
// This cell is re-evaluated and creates new tensors as we interact with the sliders to adjust parameters.
// We need to manually dispose them to avoid memory leak.
// See: https://beta.observablehq.com/@nsthorat/clean-up-deeplearn-js-tensor-generator-cells
try {
yield invalidation;
} finally {
data.x.dispose();
data.y.dispose();
}
}
// Loss: Mean Squared Error (MSE)
loss = (predictions, labels) => {
return tf.tidy(() =>
predictions.sub(labels)
.square()
.mean()
);
}
cubicPredictions = {
if (!inputData_a) return html``;
const predictions = tf.tidy(() => {
// Model parameters:
const a = tf.variable(tf.zeros([1]));
const b = tf.variable(tf.zeros([1]));
const c = tf.variable(tf.zeros([1]));
const d = tf.variable(tf.zeros([1]));
// Model: f(x) = a * x^3 + b * x^2 + c * x + d
const f = (x) =>
a.mul(x.pow(tf.scalar(3)))
.add(b.mul(x.square()))
.add(c.mul(x))
.add(d);
// Optimizer: Stochastic Gradient Descent (SGD)
const optimizer = tf.train.sgd(0.8);
// Training Loop
for (let i = 0; i < numIterations; i++) {
optimizer.minimize(() => loss(f(inputData_a.x), inputData_a.y));
}
const predTensor = f(inputData_a.x);
// tf.tidy() will not clean up variables
[a, b, c, d].forEach(x => {
x.dispose();
});
return predTensor;
});
yield predictions;
try {
yield invalidation;
} finally {
predictions.dispose();
}
}
nn = {
const model = tf.sequential({
name: 'nn',
layers: [
// Neural network with one hidden layer
tf.layers.dense({units: numNodes, activation: 'relu', inputShape: [1]}),
tf.layers.dense({units: 1})
]
});
model.compile({loss: 'meanSquaredError', optimizer: 'adam'});
return model;
}
nnTrainInfo = {
if (!inputData_b) return null;
const TRAIN_BATCHES = 50;
const BATCH_SIZE = 10;
for (let i = 0; i < TRAIN_BATCHES; i++) {
const history = await nn.fit(
inputData_b.x,
inputData_b.y,
{batchSize: BATCH_SIZE, epochs: 1}
);
const loss = history.history.loss[0];
yield {i, loss};
}
}
nnPredictions = {
if (!inputData_b) return null;
nnTrainInfo; // Just to subscribe changes in training info
const predictions = tf.tidy(() =>
nn.model.predict(inputData_b.x.expandDims(1))
);
yield predictions;
try {
yield invalidation;
} finally {
predictions.dispose();
}
}
Purpose-built for displays of data
Observable is your go-to platform for exploring data and creating expressive data visualizations. Use reactive JavaScript notebooks for prototyping and a collaborative canvas for visual data exploration and dashboard creation.
