slightly compressed the tensorflow require

README.md

@@ -90,6 +90,8 @@ cd /similarity
 npm i
 ```
 
+Add the `-t` flag to the script to use the much faster tfjs-node implementation instead of the default tfjs.
+
 Run the t-SNE script:
 ```sh
 node tsne.js -i /path/to/images

similarity/tsne.js

@@ -20,57 +20,50 @@ console.log('starting with', process.argv);
 
 const inputPath = argv.i;
 const inputFormat = argv.f || 'jpg';
-const useTfjsnode = argv.t || false;
+const tfVersion = argv.t ? 'tfjs-node' : 'tfjs'
+const tf = require("@tensorflow/" + tfVersion);
 
-var tf
-if(useTfjsnode) {
-  tf = require('@tensorflow/tfjs-node');
-}
-else {
-  tf = require('@tensorflow/tfjs');
-}
-
-const saveCsv = async (data, filename) => { 
-  const csv = d3.csvFormat(data) 
+const saveCsv = async (data, filename) => {
+  const csv = d3.csvFormat(data)
   fs.writeFileSync(filename, csv);
 }
 
 async function run() {
   // todo load via tf
   // const MOBILENET_MODEL_PATH ='https://storage.googleapis.com/tfjs-models/tfjs/mobilenet_v1_0.25_224/model.json';
   // mobilenet = await tf.loadLayersModel(MOBILENET_MODEL_PATH);
- mobilenet = await require('@tensorflow-models/mobilenet').load()
-
- const files = await glob(inputPath + '/*.' + inputFormat)
- console.log("found files", files.length)
- const subset = files.filter((d,i) => i < 100)
-
- const activations = await getActivations(files)
- console.log("done activations", activations.length)
- // saveJson(activations)
- // let activations = require("./activations.json")
- 
- const tsne = makeTsne(activations)
- saveCsv(tsne, "tsneRaw.csv")
- const tsneSpaced = giveSpace(tsne)
- // console.log(tsneSpaced)
- saveCsv(tsneSpaced, "tsne.csv")
- console.log("done")
+  mobilenet = await require('@tensorflow-models/mobilenet').load()
+
+  const files = await glob(inputPath + '/*.' + inputFormat)
+  console.log("found files", files.length)
+  const subset = files.filter((d, i) => i < 100)
+
+  const activations = await getActivations(files)
+  console.log("done activations", activations.length)
+  // saveJson(activations)
+  // let activations = require("./activations.json")
+
+  const tsne = makeTsne(activations)
+  saveCsv(tsne, "tsneRaw.csv")
+  const tsneSpaced = giveSpace(tsne)
+  // console.log(tsneSpaced)
+  saveCsv(tsneSpaced, "tsne.csv")
+  console.log("done")
 }
 
-function saveJson(data){
+function saveJson(data) {
   fs.writeFileSync("activations.json", JSON.stringify(data), "utf8")
 }
 
-function giveSpace(nodes){
+function giveSpace(nodes) {
   console.log("running space distribution")
   const simulation = d3.forceSimulation(nodes)
-      // .force("charge", d3.forceManyBody().strength(-0.0000002).distanceMin(0.001))
-       .force("charge", d3.forceManyBody().strength(-0.000001).distanceMin(0.001))
-        .force("center", d3.forceCenter(0.5,0.5))
-        //.on("tick", ()=> {  });
-        .stop()
-    
+    // .force("charge", d3.forceManyBody().strength(-0.0000002).distanceMin(0.001))
+    .force("charge", d3.forceManyBody().strength(-0.000001).distanceMin(0.001))
+    .force("center", d3.forceCenter(0.5, 0.5))
+    //.on("tick", ()=> {  });
+    .stop()
+
   //for (var i = 0, n = Math.ceil(Math.log(simulation.alphaMin()) / Math.log(1 - simulation.alphaDecay())); i < n; ++i) {
   for (var i = 0; i < 100; ++i) {
     console.log("space tick", i)
@@ -87,7 +80,7 @@ function giveSpace(nodes){
 }
 
 
-function makeTsne(activations){
+function makeTsne(activations) {
   const opt = {
     epsilon: 10,    // epsilon is learning rate (10 = default)
     perplexity: 30, // roughly how many neighbors each point influences (30 = default)
@@ -100,47 +93,47 @@ function makeTsne(activations){
 
   tsne.initDataRaw(activations.map(d => d.activation));
 
-  for(var k = 0; k < 2000; k++) {
+  for (var k = 0; k < 2000; k++) {
     console.log("tsne tick", k)
     tsne.step()
   }
 
   const output = tsne.getSolution();
 
   var xExtent = d3.extent(output, function (d) {
-      return d[0]
+    return d[0]
   })
   var yExtent = d3.extent(output, function (d) {
-      return d[1]
+    return d[1]
   })
 
   var x = d3.scaleLinear().range([0, 1]).domain(xExtent)
   var y = d3.scaleLinear().range([0, 1]).domain(yExtent)
 
   const outputScaled = output.map(function (d) {
-      return [x(d[0]), y(d[1])]
+    return [x(d[0]), y(d[1])]
   })
 
-  const merged = activations.map((d,i) => {
-   const coordinate = outputScaled[i]
-   return {
-     id: d.id,
-     x: coordinate[0],
-     y: coordinate[1]
-   }
+  const merged = activations.map((d, i) => {
+    const coordinate = outputScaled[i]
+    return {
+      id: d.id,
+      x: coordinate[0],
+      y: coordinate[1]
+    }
   })
 
   return merged
 }
 
 
-async function getActivations(files){
+async function getActivations(files) {
   const pool = []
-  for(let file of files){
+  for (let file of files) {
     const activation = await getActivation(file)
     const id = path.basename(file, '.' + inputFormat);
-    if(activation){
-      pool.push({id,activation})
+    if (activation) {
+      pool.push({ id, activation })
       console.log(file)
     } else {
       console.log("error with", file)
@@ -160,7 +153,7 @@ async function getActivation(file) {
     preds.dispose();
 
     return result
-  } catch (e){
+  } catch (e) {
     console.log(e)
     return null
   }

similarity/umap.js

@@ -8,7 +8,8 @@ var path = require("path");
 const glob = require('glob-promise');
 const localPath = i => path.relative(process.cwd(), i)
 const argv = require('minimist')(process.argv.slice(2));
-const tf = require('@tensorflow/tfjs-node');
+const tfVersion = argv.t ? 'tfjs-node' : 'tfjs'
+const tf = require("@tensorflow/" + tfVersion);
 const UMAP = require("umap-js").UMAP
 const {
   createCanvas,