[visualize-convnet] Add new example: visualize-convnet

visualize-convnet/.babelrc

@@ -0,0 +1,18 @@
+{
+  "presets": [
+    [
+      "env",
+      {
+        "esmodules": false,
+        "targets": {
+          "browsers": [
+            "> 3%"
+          ]
+        }
+      }
+    ]
+  ],
+  "plugins": [
+    "transform-runtime"
+  ]
+}

visualize-convnet/.gitignore

@@ -0,0 +1,4 @@
+*.png
+model.json
+group*shard*
+*.h5

visualize-convnet/README.md

@@ -0,0 +1,66 @@
+# TensorFlow.js Example: Visualizing Convnet Filters
+
+## Description
+
+This TensorFlow.js example demonstrates some techniques of visualizing
+the internal workings of a convolutional neural network (convnet), including:
+
+- Finding what convolutional layers' filters are sensitive to after
+  training: calculating maximally-activating input image for
+  convolutional filters through gradient ascent in the input space.
+- Getting the internal activation of a convnet by uisng the
+  functional model API of TensorFlow.js
+- Finding which part of an input image is most relevant to the
+  classification decision made by a convnet (VGG16 in this case),
+  using the gradient-based class activation map (CAM) approach.
+
+## How to use this demo
+
+Run the command:
+```sh
+yarn visualize
+```
+
+This will automatically
+
+1. install the necessary Python dependencies. If the required
+   Python package (keras, tensorflow and tensorflowjs) are already installed,
+   this step will be a no-op. However, to prevent this step from
+   modifying your global Python environment, you may run this demo from
+   a [virtualenv](https://virtualenv.pypa.io/en/latest/) or
+   [pipenv](https://pipenv.readthedocs.io/en/latest/).
+2. download and convert the VGG16 model to TensorFlow.js format
+3. launch a Node.js script to load the converted model and compute
+   the maximally-activating input images for the convnet's filters
+   using gradient ascent in the input space and save them as image
+   files under the `dist/filters` directory
+4. launch a Node.js script to calculate the internal convolutional
+   layers' activations and th gradient-based class activation
+   map (CAM) and save them as image files under the
+   `dist/activation` directory.
+5. compile and launch the web view using parcel
+
+Step 3 and 4 (especially step 3) involve relatively heavy computation
+and is best done usnig tfjs-node-gpu instead of the default
+tfjs-node. This requires that a CUDA-enabled GPU and the necessary
+driver and libraries are installed on your system.
+
+Assuming those prerequisites are met, do:
+
+```sh
+yarn visualize --gpu
+```
+
+You may also increase the number of filters to visualize per convolutional
+layer from the default 8 to a larger value, e.g., 32:
+
+```sh
+yarn visualize --gpu --filters 32
+```
+
+The default image used for the internal-activation and CAM visualization is
+"owl.jpg". You can switch to another image by using the "--image" flag, e.g.,
+
+```sh
+yarn visualize --image dog.jpg
+```

visualize-convnet/cam.js

@@ -0,0 +1,129 @@
+/**
+ * @license
+ * Copyright 2019 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+
+/**
+ * This script contains a function that performs the following operations:
+ *
+ * Get visual interpretation of which parts of the image more most
+ *    responsible for a convnet's classification decision, using the
+ *    gradient-based class activation map (CAM) method.
+ *    See function `gradClassActivationMap()`.
+ */
+
+const tf = require('@tensorflow/tfjs');
+const utils = require('./utils');
+
+/**
+ * Calculate class activation map (CAM) and overlay it on input image.
+ *
+ * This function automatically finds the last convolutional layer, get its
+ * output (activation) under the input image, weights its filters by the
+ * gradient of the class output with respect to them, and then collapses along
+ * the filter dimension.
+ *
+ * @param {tf.Sequential} model A TensorFlow.js sequential model, assumed to
+ *   contain at least one convolutional layer.
+ * @param {number} classIndex Index to class in the model's final classification
+ *   output.
+ * @param {tf.Tensor4d} x Input image, assumed to have shape
+ *   `[1, height, width, 3]`.
+ * @param {number} overlayFactor Optional overlay factor.
+ * @returns The input image with a heat-map representation of the class
+ *   activation map overlaid on top of it, as float32-type `tf.Tensor4d` of
+ *   shape `[1, height, width, 3]`.
+ */
+function gradClassActivationMap(model, classIndex, x, overlayFactor = 2.0) {
+  // Try to locate the last conv layer of the model.
+  let layerIndex = model.layers.length - 1;
+  while (layerIndex >= 0) {
+    if (model.layers[layerIndex].getClassName().startsWith('Conv')) {
+      break;
+    }
+    layerIndex--;
+  }
+  tf.util.assert(
+      layerIndex >= 0, `Failed to find a convolutional layer in model`);
+
+  const lastConvLayer = model.layers[layerIndex];
+  console.log(
+      `Located last convolutional layer of the model at ` +
+      `index ${layerIndex}: layer type = ${lastConvLayer.getClassName()}; ` +
+      `layer name = ${lastConvLayer.name}`);
+
+  // Get "sub-model 1", which goes from the original input to the output
+  // of the last convolutional layer.
+  const lastConvLayerOutput = lastConvLayer.output;
+  const subModel1 =
+      tf.model({inputs: model.inputs, outputs: lastConvLayerOutput});
+
+  // Get "sub-model 2", which goes from the output of the last convolutional
+  // layer to the original output.
+  const newInput = tf.input({shape: lastConvLayerOutput.shape.slice(1)});
+  layerIndex++;
+  let y = newInput;
+  while (layerIndex < model.layers.length) {
+    y = model.layers[layerIndex++].apply(y);
+  }
+  const subModel2 = tf.model({inputs: newInput, outputs: y});
+
+  return tf.tidy(() => {
+    // This function runs sub-model 2 and extracts the slice of the probability
+    // output that corresponds to the desired class.
+    const convOutput2ClassOutput = (input) =>
+        subModel2.apply(input, {training: true}).gather([classIndex], 1);
+    // This is the gradient function of the output corresponding to the desired
+    // class with respect to its input (i.e., the output of the last
+    // convolutional layer of the original model).
+    const gradFunction = tf.grad(convOutput2ClassOutput);
+
+    // Calculate the values of the last conv layer's output.
+    const lastConvLayerOutputValues = subModel1.apply(x);
+    // Calculate the values of gradients of the class output w.r.t. the output
+    // of the last convolutional layer.
+    const gradValues = gradFunction(lastConvLayerOutputValues);
+
+    // Pool the gradient values within each filter of the last convolutional
+    // layer, resulting in a tensor of shape [numFilters].
+    const pooledGradValues = tf.mean(gradValues, [0, 1, 2]);
+    // Scale the convlutional layer's output by the pooled gradients, using
+    // broadcasting.
+    const scaledConvOutputValues =
+        lastConvLayerOutputValues.mul(pooledGradValues);
+
+    // Create heat map by averaging and collapsing over all filters.
+    let heatMap = scaledConvOutputValues.mean(-1);
+
+    // Discard negative values from the heat map and normalize it to the [0, 1]
+    // interval.
+    heatMap = heatMap.relu();
+    heatMap = heatMap.div(heatMap.max()).expandDims(-1);
+
+    // Up-sample the heat map to the size of the input image.
+    heatMap = tf.image.resizeBilinear(heatMap, [x.shape[1], x.shape[2]]);
+
+    // Apply an RGB colormap on the heatMap. This step is necessary because
+    // the heatMap is a 1-channel (grayscale) image. It needs to be converted
+    // into a color (RGB) one through this function call.
+    heatMap = utils.applyColorMap(heatMap);
+
+    // To form the final output, overlay the color heat map on the input image.
+    heatMap = heatMap.mul(overlayFactor).add(x.div(255));
+    return heatMap.div(heatMap.max()).mul(255);
+  });
+}
+
+module.exports = {gradClassActivationMap};

visualize-convnet/filters.js

@@ -0,0 +1,174 @@
+/**
+ * @license
+ * Copyright 2019 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+
+/**
+ * Algorithms for analyzing and visualizing the convolutional filters
+ * internal to a convnet.
+ *
+ * 1. Retrieving internal activations of a convnet.
+ *    See function `writeInternalActivationAndGetOutput()`.
+ * 2. Calculate maximally-activating input image for convnet filters, using
+ *    gradient ascent in input space.
+ *    See function `inputGradientAscent()`.
+ **/
+
+const path = require('path');
+const tf = require('@tensorflow/tfjs');
+const utils = require('./utils');
+
+/**
+ * Write internal activation of conv layers to file; Get model output.
+ *
+ * @param {tf.Model} model The model of interest.
+ * @param {string[]} layerNames Names of layers of interest.
+ * @param {tf.Tensor4d} inputImage The input image represented as a 4D tensor
+ *   of shape [1, height, width, 3].
+ * @param {number} numFilters Number of filters to run for each convolutional
+ *   layer. If it exceeds the number of filters of a convolutional layer, it
+ *   will be cut off.
+ * @param {string} outputDir Path to the directory to which the image files
+ *   representing the activation will be saved.
+ * @return modelOutput: final output of the model as a tf.Tensor.
+ *         layerName2FilePaths: an object mapping layer name to the paths to the
+ *           image files saved for the layer's activation.
+ *         layerName2FilePaths: an object mapping layer name to the height
+ *           and width of the layer's filter outputs.
+ */
+async function writeInternalActivationAndGetOutput(
+    model, layerNames, inputImage, numFilters, outputDir) {
+  const layerName2FilePaths = {};
+  const layerName2ImageDims = {};
+  const layerOutputs =
+      layerNames.map(layerName => model.getLayer(layerName).output);
+
+  // Construct a model that returns all the desired internal activations,
+  // in addition to the final output of the original model.
+  const compositeModel = tf.model(
+      {inputs: model.input, outputs: layerOutputs.concat(model.outputs[0])});
+
+  // `outputs` is an array of `tf.Tensor`s consisting of the internal-activation
+  // values and the final output value.
+  const outputs = compositeModel.predict(inputImage);
+
+  for (let i = 0; i < outputs.length - 1; ++i) {
+    const layerName = layerNames[i];
+    // Split the activation of the convolutional layer by filter.
+    const activationTensors =
+        tf.split(outputs[i], outputs[i].shape[outputs[i].shape.length - 1], -1);
+    const actualNumFilters = numFilters <= activationTensors.length ?
+        numFilters :
+        activationTensors.length;
+    const filePaths = [];
+    let imageTensorShape;
+    for (let j = 0; j < actualNumFilters; ++j) {
+      // Format activation tensors and write them to disk.
+      const imageTensor = tf.tidy(
+          () => deprocessImage(tf.tile(activationTensors[j], [1, 1, 1, 3])));
+      const outputFilePath = path.join(outputDir, `${layerName}_${j + 1}.png`);
+      filePaths.push(outputFilePath);
+      await utils.writeImageTensorToFile(imageTensor, outputFilePath);
+      imageTensorShape = imageTensor.shape;
+    }
+    layerName2FilePaths[layerName] = filePaths;
+    layerName2ImageDims[layerName] = imageTensorShape.slice(1, 3);
+    tf.dispose(activationTensors);
+  }
+  tf.dispose(outputs.slice(0, outputs.length - 1));
+  return {
+    modelOutput: outputs[outputs.length - 1],
+    layerName2FilePaths,
+    layerName2ImageDims
+  };
+}
+
+
+/**
+ * Generate the maximally-activating input image for a conv2d layer filter.
+ *
+ * Uses gradient ascent in input space.
+ *
+ * @param {tf.Model} model The model that the convolutional layer of interest
+ *   belongs to.
+ * @param {string} layerName Name of the convolutional layer.
+ * @param {number} filterIndex Index to the filter of interest. Must be
+ *   < number of filters of the conv2d layer.
+ * @param {number} iterations Number of gradient-ascent iterations.
+ * @return {tf.Tensor} The maximally-activating input image as a tensor.
+ */
+function inputGradientAscent(model, layerName, filterIndex, iterations = 40) {
+  return tf.tidy(() => {
+    const imageH = model.inputs[0].shape[1];
+    const imageW = model.inputs[0].shape[2];
+    const imageDepth = model.inputs[0].shape[3];
+
+    // Create an auxiliary model of which input is the same as the original
+    // model but the output is the output of the convolutional layer of
+    // interest.
+    const layerOutput = model.getLayer(layerName).output;
+    const auxModel = tf.model({inputs: model.inputs, outputs: layerOutput});
+
+    // This function calculates the value of the convolutional layer's
+    // output at the designated filter index.
+    const lossFunction = (input) =>
+        auxModel.apply(input, {training: true}).gather([filterIndex], 3);
+
+    // This function (`gradient`) calculates the gradient of the convolutional
+    // filter's output with respect to the input image.
+    const gradients = tf.grad(lossFunction);
+
+    // Form a random image as the starting point of the gradient ascent.
+    let image = tf.randomUniform([1, imageH, imageW, imageDepth], 0, 1)
+                    .mul(20)
+                    .add(128);
+
+    for (let i = 0; i < iterations; ++i) {
+      const scaledGrads = tf.tidy(() => {
+        const grads = gradients(image);
+        const norm =
+            tf.sqrt(tf.mean(tf.square(grads))).add(tf.ENV.get('EPSILON'));
+        // Important trick: scale the gradient with the magnitude (norm)
+        // of the gradient.
+        return grads.div(norm);
+      });
+      // Perform one step of gradient ascent: Update the image along the
+      // direction of the gradient.
+      image = tf.clipByValue(image.add(scaledGrads), 0, 255);
+    }
+    return deprocessImage(image);
+  });
+}
+
+/** Center and scale input image so the pixel values fall into [0, 255]. */
+function deprocessImage(x) {
+  return tf.tidy(() => {
+    const {mean, variance} = tf.moments(x);
+    x = x.sub(mean);
+    // Add a small positive number (EPSILON) to the denominator to prevent
+    // division-by-zero.
+    x = x.div(tf.sqrt(variance).add(tf.ENV.get('EPSILON')));
+    // Clip to [0, 1].
+    x = x.add(0.5);
+    x = tf.clipByValue(x, 0, 1);
+    x = x.mul(255);
+    return tf.clipByValue(x, 0, 255).asType('int32');
+  });
+}
+
+module.exports = {
+  inputGradientAscent,
+  writeInternalActivationAndGetOutput
+};