|
| 1 | +/** |
| 2 | + * @license |
| 3 | + * Copyright 2019 Google LLC. All Rights Reserved. |
| 4 | + * Licensed under the Apache License, Version 2.0 (the "License"); |
| 5 | + * you may not use this file except in compliance with the License. |
| 6 | + * You may obtain a copy of the License at |
| 7 | + * |
| 8 | + * http://www.apache.org/licenses/LICENSE-2.0 |
| 9 | + * |
| 10 | + * Unless required by applicable law or agreed to in writing, software |
| 11 | + * distributed under the License is distributed on an "AS IS" BASIS, |
| 12 | + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 13 | + * See the License for the specific language governing permissions and |
| 14 | + * limitations under the License. |
| 15 | + * ============================================================================= |
| 16 | + */ |
| 17 | + |
| 18 | +/** |
| 19 | + * Algorithms for analyzing and visualizing the convolutional filters |
| 20 | + * internal to a convnet. |
| 21 | + * |
| 22 | + * 1. Retrieving internal activations of a convnet. |
| 23 | + * See function `writeInternalActivationAndGetOutput()`. |
| 24 | + * 2. Calculate maximally-activating input image for convnet filters, using |
| 25 | + * gradient ascent in input space. |
| 26 | + * See function `inputGradientAscent()`. |
| 27 | + **/ |
| 28 | + |
| 29 | +const path = require('path'); |
| 30 | +const tf = require('@tensorflow/tfjs'); |
| 31 | +const utils = require('./utils'); |
| 32 | + |
| 33 | +/** |
| 34 | + * Write internal activation of conv layers to file; Get model output. |
| 35 | + * |
| 36 | + * @param {tf.Model} model The model of interest. |
| 37 | + * @param {string[]} layerNames Names of layers of interest. |
| 38 | + * @param {tf.Tensor4d} inputImage The input image represented as a 4D tensor |
| 39 | + * of shape [1, height, width, 3]. |
| 40 | + * @param {number} numFilters Number of filters to run for each convolutional |
| 41 | + * layer. If it exceeds the number of filters of a convolutional layer, it |
| 42 | + * will be cut off. |
| 43 | + * @param {string} outputDir Path to the directory to which the image files |
| 44 | + * representing the activation will be saved. |
| 45 | + * @return modelOutput: final output of the model as a tf.Tensor. |
| 46 | + * layerName2FilePaths: an object mapping layer name to the paths to the |
| 47 | + * image files saved for the layer's activation. |
| 48 | + * layerName2FilePaths: an object mapping layer name to the height |
| 49 | + * and width of the layer's filter outputs. |
| 50 | + */ |
| 51 | +async function writeInternalActivationAndGetOutput( |
| 52 | + model, layerNames, inputImage, numFilters, outputDir) { |
| 53 | + const layerName2FilePaths = {}; |
| 54 | + const layerName2ImageDims = {}; |
| 55 | + const layerOutputs = |
| 56 | + layerNames.map(layerName => model.getLayer(layerName).output); |
| 57 | + |
| 58 | + // Construct a model that returns all the desired internal activations, |
| 59 | + // in addition to the final output of the original model. |
| 60 | + const compositeModel = tf.model( |
| 61 | + {inputs: model.input, outputs: layerOutputs.concat(model.outputs[0])}); |
| 62 | + |
| 63 | + // `outputs` is an array of `tf.Tensor`s consisting of the internal-activation |
| 64 | + // values and the final output value. |
| 65 | + const outputs = compositeModel.predict(inputImage); |
| 66 | + |
| 67 | + for (let i = 0; i < outputs.length - 1; ++i) { |
| 68 | + const layerName = layerNames[i]; |
| 69 | + // Split the activation of the convolutional layer by filter. |
| 70 | + const activationTensors = |
| 71 | + tf.split(outputs[i], outputs[i].shape[outputs[i].shape.length - 1], -1); |
| 72 | + const actualNumFilters = numFilters <= activationTensors.length ? |
| 73 | + numFilters : |
| 74 | + activationTensors.length; |
| 75 | + const filePaths = []; |
| 76 | + let imageTensorShape; |
| 77 | + for (let j = 0; j < actualNumFilters; ++j) { |
| 78 | + // Format activation tensors and write them to disk. |
| 79 | + const imageTensor = tf.tidy( |
| 80 | + () => deprocessImage(tf.tile(activationTensors[j], [1, 1, 1, 3]))); |
| 81 | + const outputFilePath = path.join(outputDir, `${layerName}_${j + 1}.png`); |
| 82 | + filePaths.push(outputFilePath); |
| 83 | + await utils.writeImageTensorToFile(imageTensor, outputFilePath); |
| 84 | + imageTensorShape = imageTensor.shape; |
| 85 | + } |
| 86 | + layerName2FilePaths[layerName] = filePaths; |
| 87 | + layerName2ImageDims[layerName] = imageTensorShape.slice(1, 3); |
| 88 | + tf.dispose(activationTensors); |
| 89 | + } |
| 90 | + tf.dispose(outputs.slice(0, outputs.length - 1)); |
| 91 | + return { |
| 92 | + modelOutput: outputs[outputs.length - 1], |
| 93 | + layerName2FilePaths, |
| 94 | + layerName2ImageDims |
| 95 | + }; |
| 96 | +} |
| 97 | + |
| 98 | + |
| 99 | +/** |
| 100 | + * Generate the maximally-activating input image for a conv2d layer filter. |
| 101 | + * |
| 102 | + * Uses gradient ascent in input space. |
| 103 | + * |
| 104 | + * @param {tf.Model} model The model that the convolutional layer of interest |
| 105 | + * belongs to. |
| 106 | + * @param {string} layerName Name of the convolutional layer. |
| 107 | + * @param {number} filterIndex Index to the filter of interest. Must be |
| 108 | + * < number of filters of the conv2d layer. |
| 109 | + * @param {number} iterations Number of gradient-ascent iterations. |
| 110 | + * @return {tf.Tensor} The maximally-activating input image as a tensor. |
| 111 | + */ |
| 112 | +function inputGradientAscent(model, layerName, filterIndex, iterations = 40) { |
| 113 | + return tf.tidy(() => { |
| 114 | + const imageH = model.inputs[0].shape[1]; |
| 115 | + const imageW = model.inputs[0].shape[2]; |
| 116 | + const imageDepth = model.inputs[0].shape[3]; |
| 117 | + |
| 118 | + // Create an auxiliary model of which input is the same as the original |
| 119 | + // model but the output is the output of the convolutional layer of |
| 120 | + // interest. |
| 121 | + const layerOutput = model.getLayer(layerName).output; |
| 122 | + const auxModel = tf.model({inputs: model.inputs, outputs: layerOutput}); |
| 123 | + |
| 124 | + // This function calculates the value of the convolutional layer's |
| 125 | + // output at the designated filter index. |
| 126 | + const lossFunction = (input) => |
| 127 | + auxModel.apply(input, {training: true}).gather([filterIndex], 3); |
| 128 | + |
| 129 | + // This function (`gradient`) calculates the gradient of the convolutional |
| 130 | + // filter's output with respect to the input image. |
| 131 | + const gradients = tf.grad(lossFunction); |
| 132 | + |
| 133 | + // Form a random image as the starting point of the gradient ascent. |
| 134 | + let image = tf.randomUniform([1, imageH, imageW, imageDepth], 0, 1) |
| 135 | + .mul(20) |
| 136 | + .add(128); |
| 137 | + |
| 138 | + for (let i = 0; i < iterations; ++i) { |
| 139 | + const scaledGrads = tf.tidy(() => { |
| 140 | + const grads = gradients(image); |
| 141 | + const norm = |
| 142 | + tf.sqrt(tf.mean(tf.square(grads))).add(tf.ENV.get('EPSILON')); |
| 143 | + // Important trick: scale the gradient with the magnitude (norm) |
| 144 | + // of the gradient. |
| 145 | + return grads.div(norm); |
| 146 | + }); |
| 147 | + // Perform one step of gradient ascent: Update the image along the |
| 148 | + // direction of the gradient. |
| 149 | + image = tf.clipByValue(image.add(scaledGrads), 0, 255); |
| 150 | + } |
| 151 | + return deprocessImage(image); |
| 152 | + }); |
| 153 | +} |
| 154 | + |
| 155 | +/** Center and scale input image so the pixel values fall into [0, 255]. */ |
| 156 | +function deprocessImage(x) { |
| 157 | + return tf.tidy(() => { |
| 158 | + const {mean, variance} = tf.moments(x); |
| 159 | + x = x.sub(mean); |
| 160 | + // Add a small positive number (EPSILON) to the denominator to prevent |
| 161 | + // division-by-zero. |
| 162 | + x = x.div(tf.sqrt(variance).add(tf.ENV.get('EPSILON'))); |
| 163 | + // Clip to [0, 1]. |
| 164 | + x = x.add(0.5); |
| 165 | + x = tf.clipByValue(x, 0, 1); |
| 166 | + x = x.mul(255); |
| 167 | + return tf.clipByValue(x, 0, 255).asType('int32'); |
| 168 | + }); |
| 169 | +} |
| 170 | + |
| 171 | +module.exports = { |
| 172 | + inputGradientAscent, |
| 173 | + writeInternalActivationAndGetOutput |
| 174 | +}; |
0 commit comments