async function showPredictions() {
const testExamples = 100;
const batch = data.nextTestBatch(testExamples);
// Code wrapped in a tf.tidy() function callback will have their tensors freed
// from GPU memory after execution without having to call dispose().
// The tf.tidy callback runs synchronously.
tf.tidy(() => {
const output = model.predict(batch.xs.reshape([-1, 28, 28, 1]));
// tf.argMax() returns the indices of the maximum values in the tensor along
// a specific axis. Categorical classification tasks like this one often
// represent classes as one-hot vectors. One-hot vectors are 1D vectors with
// one element for each output class. All values in the vector are 0
// except for one, which has a value of 1 (e.g. [0, 0, 0, 1, 0]). The
// output from model.predict() will be a probability distribution, so we use
// argMax to get the index of the vector element that has the highest
// probability. This is our prediction.
// (e.g. argmax([0.07, 0.1, 0.03, 0.75, 0.05]) == 3)
// dataSync() synchronously downloads the tf.tensor values from the GPU so
// that we can use them in our normal CPU JavaScript code
// (for a non-blocking version of this function, use data()).
const axis = 1;
const labels = Array.from(batch.labels.argMax(axis).dataSync());
const predictions = Array.from(output.argMax(axis).dataSync());
ui.showTestResults(batch, predictions, labels);
});
}
//
recordCost() {
const cost = tf.tidy(() => {
const guesses = this.features.matMul(this.weights).softmax();
const termOne = this.labels
.transpose()
.matMul(
guesses
.add(1e-7) // add to avoid log(0)
.log()
);
const termTwo = this.labels
.mul(-1)
.add(1)
.transpose()
.matMul(
guesses
.mul(-1)
.add(1)
.add(1e-7) // add to avoid log(0)
.log()
);
return termOne.add(termTwo)
.div(this.features.shape[0])
.mul(-1)
.get(0, 0);
});
this.costHistory.unshift(cost);
}
/**
* Given an image element, makes a prediction through mobilenet returning the
* probabilities of the top K classes.
*/
async function predict(imgElement) {
status('Predicting...');
const startTime = performance.now();
const logits = tf.tidy(() => {
// tf.fromPixels() returns a Tensor from an image element.
const img = tf.fromPixels(imgElement).toFloat();
const offset = tf.scalar(127.5);
// Normalize the image from [0, 255] to [-1, 1].
const normalized = img.sub(offset).div(offset);
// Reshape to a single-element batch so we can pass it to predict.
const batched = normalized.reshape([1, IMAGE_SIZE, IMAGE_SIZE, 3]);
// Make a prediction through mobilenet.
return mobilenet.predict(batched);
});
// Convert logits to probabilities and class names.
const classes = await getTopKClasses(logits, TOPK_PREDICTIONS);
const totalTime = performance.now() - startTime;
status(`Done in ${Math.floor(totalTime)}ms`);
// Show the classes in the DOM.
showResults(imgElement, classes);
}
async loadModel() {
this.mobilenet = await tf.loadModel(this.modelPath);
const layer = this.mobilenet.getLayer('conv_pw_13_relu');
if (this.videoReady && this.video) {
tf.tidy(() => this.mobilenet.predict(imgToTensor(this.video))); // Warm up
}
return tf.model({ inputs: this.mobilenet.inputs, outputs: layer.output });
}
async function train() {
ui.isTraining();
// We'll keep a buffer of loss and accuracy values over time.
const lossValues = [];
const accuracyValues = [];
// Iteratively train our model on mini-batches of data.
for (let i = 0; i < TRAIN_BATCHES; i++) {
const [batch, validationData] = tf.tidy(() => {
const batch = data.nextTrainBatch(BATCH_SIZE);
batch.xs = batch.xs.reshape([BATCH_SIZE, 28, 28, 1]);
let validationData;
// Every few batches test the accuracy of the model.
if (i % TEST_ITERATION_FREQUENCY === 0) {
const testBatch = data.nextTestBatch(TEST_BATCH_SIZE);
validationData = [
// Reshape the training data from [64, 28x28] to [64, 28, 28, 1] so
// that we can feed it to our convolutional neural net.
testBatch.xs.reshape([TEST_BATCH_SIZE, 28, 28, 1]), testBatch.labels
];
}
return [batch, validationData];
});
// The entire dataset doesn't fit into memory so we call train repeatedly
// with batches using the fit() method.
const history = await model.fit(
batch.xs, batch.labels,
{batchSize: BATCH_SIZE, validationData, epochs: 1});
const loss = history.history.loss[0];
const accuracy = history.history.acc[0];
// Plot loss / accuracy.
lossValues.push({'batch': i, 'loss': loss, 'set': 'train'});
ui.plotLosses(lossValues);
if (validationData != null) {
accuracyValues.push({'batch': i, 'accuracy': accuracy, 'set': 'train'});
ui.plotAccuracies(accuracyValues);
}
// Call dispose on the training/test tensors to free their GPU memory.
tf.dispose([batch, validationData]);
// tf.nextFrame() returns a promise that resolves at the next call to
// requestAnimationFrame(). By awaiting this promise we keep our model
// training from blocking the main UI thread and freezing the browser.
await tf.nextFrame();
}
}
// Add an image to retrain
addImage(labelOrInput, labelOrCallback, cb) {
let imgToAdd;
let label;
let callback;
if (labelOrInput instanceof HTMLImageElement || labelOrInput instanceof HTMLVideoElement) {
imgToAdd = labelOrInput;
label = labelOrCallback;
callback = cb;
} else {
imgToAdd = this.video;
label = labelOrInput;
callback = labelOrCallback;
}
if (typeof label === 'string') {
if (!this.mapStringToIndex.includes(label)) {
label = this.mapStringToIndex.push(label) - 1;
} else {
label = this.mapStringToIndex.indexOf(label);
}
}
if (this.modelLoaded) {
tf.tidy(() => {
const processedImg = imgToTensor(imgToAdd);
const prediction = this.mobilenetModified.predict(processedImg);
const y = tf.tidy(() => tf.oneHot(tf.tensor1d([label], 'int32'), this.numClasses));
if (this.xs == null) {
this.xs = tf.keep(prediction);
this.ys = tf.keep(y);
this.hasAnyTrainedClass = true;
} else {
const oldX = this.xs;
this.xs = tf.keep(oldX.concat(prediction, 0));
const oldY = this.ys;
this.ys = tf.keep(oldY.concat(y, 0));
oldX.dispose();
oldY.dispose();
y.dispose();
}
});
if (callback) {
callback();
}
}
}
tf.tidy(() => {
const processedImg = imgToTensor(imgToAdd);
const prediction = this.mobilenetModified.predict(processedImg);
const y = tf.tidy(() => tf.oneHot(tf.tensor1d([label], 'int32'), this.numClasses));
if (this.xs == null) {
this.xs = tf.keep(prediction);
this.ys = tf.keep(y);
this.hasAnyTrainedClass = true;
} else {
const oldX = this.xs;
this.xs = tf.keep(oldX.concat(prediction, 0));
const oldY = this.ys;
this.ys = tf.keep(oldY.concat(y, 0));
oldX.dispose();
oldY.dispose();
y.dispose();
}
});
//
train() {
const { batchSize } = this.options;
const batchCount = Math.floor(
this.features.shape[0] / batchSize
);
for (let i = 0; i < this.options.iterations ; i++) {
this.recordWeights();
for (let batch = 0; batch < batchCount ; batch++) {
const startIndex = batch * batchSize;
this.weights = tf.tidy(() => {
const featureSlice = this.features.slice([startIndex, 0],[batchSize, -1]);
const labelsSlice = this.labels.slice([startIndex, 0],[batchSize, -1]);
return this.gradientDescent(featureSlice, labelsSlice);
});
}
this.recordCost();
this.updateLearningRate();
}
this.costHistory.reverse();
}
transfer(input) {
const image = tf.fromPixels(input);
const result = tf.tidy(() => {
const conv1 = this.convLayer(image, 1, true, 0);
const conv2 = this.convLayer(conv1, 2, true, 3);
const conv3 = this.convLayer(conv2, 2, true, 6);
const res1 = this.residualBlock(conv3, 9);
const res2 = this.residualBlock(res1, 15);
const res3 = this.residualBlock(res2, 21);
const res4 = this.residualBlock(res3, 27);
const res5 = this.residualBlock(res4, 33);
const convT1 = this.convTransposeLayer(res5, 64, 2, 39);
const convT2 = this.convTransposeLayer(convT1, 32, 2, 42);
const convT3 = this.convLayer(convT2, 1, false, 45);
const outTanh = tf.tanh(convT3);
const scaled = tf.mul(this.timesScalar, outTanh);
const shifted = tf.add(this.plusScalar, scaled);
const clamped = tf.clipByValue(shifted, 0, 255);
const normalized = tf.div(clamped, tf.scalar(255.0));
return normalized;
});
return array3DToImage(result);
}
/**
* Adds an example to the controller dataset.
* @param {Tensor} example A tensor representing the example. It can be an image,
* an activation, or any other type of Tensor.
* @param {number} label The label of the example. Should be an umber.
*/
addExample(example, label) {
// One-hot encode the label.
const y = tf.tidy(
() => tf.oneHot(tf.tensor1d([label]).toInt(), this.numClasses));
if (this.xs == null) {
// For the first example that gets added, keep example and y so that the
// ControllerDataset owns the memory of the inputs. This makes sure that
// if addExample() is called in a tf.tidy(), these Tensors will not get
// disposed.
this.xs = tf.keep(example);
this.ys = tf.keep(y);
} else {
const oldX = this.xs;
this.xs = tf.keep(oldX.concat(example, 0));
const oldY = this.ys;
this.ys = tf.keep(oldY.concat(y, 0));
oldX.dispose();
oldY.dispose();
y.dispose();
}
}
/* linear predictor */
function predict(x) { return tf.tidy(function() { return a.mul(x).add(b); }); }
/* eslint consistent-return: 0 */
async predict(inputNumOrCallback, numOrCallback = null, cb = null) {
let imgToPredict = this.video;
let numberOfClasses = 10;
let callback;
if (typeof inputNumOrCallback === 'function') {
imgToPredict = this.video;
callback = inputNumOrCallback;
} else if (inputNumOrCallback instanceof HTMLImageElement) {
imgToPredict = inputNumOrCallback;
} else if (typeof inputNumOrCallback === 'object' && inputNumOrCallback.elt instanceof HTMLImageElement) {
imgToPredict = inputNumOrCallback.elt;
} else if (inputNumOrCallback instanceof HTMLVideoElement) {
if (!this.video) {
this.video = processVideo(inputNumOrCallback, this.imageSize);
}
imgToPredict = this.video;
} else if (typeof numOrCallback === 'number') {
imgToPredict = this.video;
numberOfClasses = inputNumOrCallback;
}
if (typeof numOrCallback === 'function') {
callback = numOrCallback;
} else if (typeof numOrCallback === 'number') {
numberOfClasses = numOrCallback;
}
if (typeof cb === 'function') {
callback = cb;
}
if (!this.modelLoaded || !this.videoReady) {
this.waitingPredictions.push({ imgToPredict, num: numberOfClasses || this.topKPredictions, callback });
} else {
// If there is no custom model, then run over the original mobilenet
if (!this.customModel) {
const logits = tf.tidy(() => {
const pixels = tf.fromPixels(imgToPredict).toFloat();
const resized = tf.image.resizeBilinear(pixels, [this.imageSize, this.imageSize]);
const offset = tf.scalar(127.5);
const normalized = resized.sub(offset).div(offset);
const batched = normalized.reshape([1, this.imageSize, this.imageSize, 3]);
return this.mobilenet.predict(batched);
});
const results = await ImageClassifier.getTopKClasses(logits, numberOfClasses || this.topKPredictions, callback);
return results;
}
this.isPredicting = true;
const predictedClass = tf.tidy(() => {
const processedImg = imgToTensor(imgToPredict);
const activation = this.mobilenetModified.predict(processedImg);
const predictions = this.customModel.predict(activation);
return predictions.as1D().argMax();
});
let classId = (await predictedClass.data())[0];
await tf.nextFrame();
if (callback) {
if (this.mapStringToIndex.length > 0) {
classId = this.mapStringToIndex[classId];
}
callback(classId);
}
}
}