[simple-object-detection] A few improvements (#224)

- Add a command line flag for switching between tfjs-node
  and tfjs-node-gpu
- Change README.md accordingly
- Add live demo link to README.md
- Remove unnecessary tensor creation, which causes problems
  because tfjs-node/tfjs-node-gpu is imported

Author: caisq

simple-object-detection/README.md

@@ -35,7 +35,7 @@ are too many training examples generated. In the meantime, having a large
 number of training examples benefits the accuracy of the model after
 training. The default number of examples is 2000. You can adjust the number
 of examples by using the `--numExamples` flag of the `yarn train` command.
-For example, the hosted model is trained with the 10000 examples, using
+For example, the hosted model is trained with the 20000 examples, using
 the command line:
 
 ```sh
@@ -52,18 +52,13 @@ See `train.js` for other adjustable parameters.
 Note that by default, the model is trained using the CPU version of tfjs-node.
 If you machine is equipped with a CUDA(R) GPU, you may switch to using
 tfjs-node-gpu, which will significantly shorten the training time. Specifically,
-in `package.json`, change the dependency `tfjs-node` to `tfjs-node-gpu`. Then,
-in `train.js`, change the line 
+add the `--gpu` flag to the command above, i.e.,
 
-```js
-require('@tensorflow/tfjs-node');
-```
-
-to
-
-```js
-require('@tensorflow/tfjs-node-gpu');
+```sh
+yarn train --gpu \
+    --numExamples 20000 \
+    --initialTransferEpochs 100 \
+    --fineTuningEpochs 200
 ```
 
-TODO(cais): Add the link below.
-[See this example live!](./README.md)
+[See this example live!](https://storage.googleapis.com/tfjs-examples/simple-object-detection/dist/index.html)

simple-object-detection/package.json

@@ -9,7 +9,7 @@
     "node": ">=8.9.0"
   },
   "dependencies": {
-    "@tensorflow/tfjs-node": "0.2.1",
+    "@tensorflow/tfjs": "0.14.2",
     "argparse": "^1.0.10",
     "canvas": "2.0.1",
     "node-fetch": "2.2.1"
@@ -22,6 +22,8 @@
     "train": "node train.js"
   },
   "devDependencies": {
+    "@tensorflow/tfjs-node": "0.2.1",
+    "@tensorflow/tfjs-node-gpu": "0.2.1",
     "babel-core": "^6.26.3",
     "babel-plugin-transform-runtime": "~6.23.0",
     "babel-polyfill": "~6.26.0",

simple-object-detection/train.js

@@ -22,18 +22,6 @@ const argparse = require('argparse');
 const canvas = require('canvas');
 const tf = require('@tensorflow/tfjs');
 const synthesizer = require('./synthetic_images');
-const fetch = require('node-fetch');
-
-// To train the model using CUDA/CuDNN,
-//   1) Make sure you have a CUDA-enabled GPU on your system.
-//   2) Install the necessary NVIDIA driver, CUDA toolkit and CuDNN library.
-//   3) Change the "@tensorflow/tfjs-node" dependency to
-//      "@tensorflow/tfjs-node-gpu" in package.json.
-//   4) Change the following line to:
-//      require('@tensorflow/tfjs-node-gpu');
-require('@tensorflow/tfjs-node');
-
-global.fetch = fetch;
 
 const CANVAS_SIZE = 224;  // Matches the input size of MobileNet.
 
@@ -47,7 +35,7 @@ const topLayerName =
 // Used to scale the first column (0-1 shape indicator) of `yTrue`
 // in order to ensure balanced contributions to the final loss value
 // from shape and bounding-box predictions.
-const LABEL_MULTIPLIER = tf.tensor1d([CANVAS_SIZE, 1, 1, 1, 1]);
+const LABEL_MULTIPLIER = [CANVAS_SIZE, 1, 1, 1, 1];
 
 /**
  * Custom loss function for object detection.
@@ -155,6 +143,10 @@ async function buildObjectDetectionModel() {
   const numLines = 10;
 
   const parser = new argparse.ArgumentParser();
+  parser.addArgument('--gpu', {
+    action: 'storeTrue',
+    help: "Use tfjs-node-gpu for training (required CUDA and CuDNN)"
+  });
   parser.addArgument(
       '--numExamples',
       {type: 'int', defaultValue: 2000, help: 'Number of training exapmles'});
@@ -181,6 +173,14 @@ async function buildObjectDetectionModel() {
   });
   const args = parser.parseArgs();
 
+  if (args.gpu) {
+    console.log('Training using GPU.');
+    require('@tensorflow/tfjs-node-gpu');
+  } else {
+    console.log('Training using CPU.');
+    require('@tensorflow/tfjs-node');
+  }
+
   const modelSaveURL = 'file://./dist/object_detection_model';
 
   const tBegin = tf.util.now();

simple-object-detection/yarn.lock

@@ -688,38 +688,51 @@
   version "1.1.0"
   resolved "https://registry.yarnpkg.com/@protobufjs/utf8/-/utf8-1.1.0.tgz#a777360b5b39a1a2e5106f8e858f2fd2d060c570"
 
-"@tensorflow/[email protected].1":
-  version "0.7.1"
-  resolved "https://registry.yarnpkg.com/@tensorflow/tfjs-converter/-/tfjs-converter-0.7.1.tgz#af33b89f34fdfbdfbe8d9361c5542d7db514364d"
-  integrity sha512-kMQqM3GI5bBl6YQ98jCJNu49NvAWGeI3iVxO3ZqOYtN90lb/+3dSBelDo2LHFXc8jnJHpFOwkFPSZlCVrVGRag==
+"@tensorflow/[email protected].2":
+  version "0.7.2"
+  resolved "https://registry.yarnpkg.com/@tensorflow/tfjs-converter/-/tfjs-converter-0.7.2.tgz#49e578f71eb82d821af05176754c3452b42cfe9c"
+  integrity sha512-m46mtaF57x2NcxlNUKdJOCUp3ZSJU9bp9MzyEQ0Iz1bW2kKIxx1DDRjuP0fAeHX5H5Mh/tWIHB9yK6NwLz+aQQ==
   dependencies:
     "@types/long" "~3.0.32"
     protobufjs "~6.8.6"
 
-"@tensorflow/[email protected].2":
-  version "0.14.2"
-  resolved "https://registry.yarnpkg.com/@tensorflow/tfjs-core/-/tfjs-core-0.14.2.tgz#b1ee6af0d893782a1c3b2c988091c36c42f1d7a4"
-  integrity sha512-VVbcu6H3ioKCkfkep/gQASfzPnQt3C5v+4ppH9pQ6Lf0lD+l3NMuMJYxa8Wjac1TfiWhFEX58bJvhpMfTGsUlg==
+"@tensorflow/[email protected].5":
+  version "0.14.5"
+  resolved "https://registry.yarnpkg.com/@tensorflow/tfjs-core/-/tfjs-core-0.14.5.tgz#17c3beeec31c4cd92b0f79a5ef30c4975a11e408"
+  integrity sha512-CSUgKuC17J1Ylr1s6iD1k2/tJr9lD16sUEjtzJbtiuTYCELOwujGK/1htunA7o3BwLuU7aqEI92MoKElEKa7qA==
   dependencies:
     "@types/seedrandom" "2.4.27"
     "@types/webgl-ext" "0.0.30"
     "@types/webgl2" "0.0.4"
     seedrandom "2.4.3"
 
-"@tensorflow/[email protected].4":
-  version "0.1.4"
-  resolved "https://registry.yarnpkg.com/@tensorflow/tfjs-data/-/tfjs-data-0.1.4.tgz#4732a2505c5d7d0d8ba0458f5b92837fd02ba442"
-  integrity sha512-YwaWNZnJj++QFEHQ1AKLqn2fvnmSp1X6CJ5YL5XJhq+m8P0AoouW9IpumCgO6WSjnD1M83/cVGZXzDIgJ4IlLg==
+"@tensorflow/[email protected].7":
+  version "0.1.7"
+  resolved "https://registry.yarnpkg.com/@tensorflow/tfjs-data/-/tfjs-data-0.1.7.tgz#8a4e43313b3d63cdfab719c0c1c47ced2ef321e3"
+  integrity sha512-RENjeBdBLq7GS9594kQx2GbM0WQV16VfxzzB0j2sq5vJh9GZQi2DB5Emq2LqZWs5rSeh7PDHZylGOn/ve6f8PA==
   dependencies:
     "@types/node-fetch" "^2.1.2"
     node-fetch "~2.1.2"
     seedrandom "~2.4.3"
-    utf8 "~2.1.2"
 
-"@tensorflow/[email protected]":
-  version "0.9.1"
-  resolved "https://registry.yarnpkg.com/@tensorflow/tfjs-layers/-/tfjs-layers-0.9.1.tgz#685b741cf5c42aa9b2621485ffafdcc43c8db699"
-  integrity sha512-TZTi59E3rGdVTJCf/AEX1arigPp0426FJLGfiKzTXp3skEuubwDM9XlwiXWcB5+l+Pjvwg4FMNXwEyajXIxX2w==
+"@tensorflow/[email protected]":
+  version "0.9.2"
+  resolved "https://registry.yarnpkg.com/@tensorflow/tfjs-layers/-/tfjs-layers-0.9.2.tgz#f5c1918d1a9660096f259cd1f99f59e689b41b69"
+  integrity sha512-peB824cEXRBy5IgZPIodd8zpQ/54VGOYbR+zY+Q1Le7v3Np05EoDcL8Z98MtpBHo6jOM7b/3Lf2zjfJVv2qxJA==
+
+"@tensorflow/[email protected]":
+  version "0.2.1"
+  resolved "https://registry.yarnpkg.com/@tensorflow/tfjs-node-gpu/-/tfjs-node-gpu-0.2.1.tgz#f8843ff3dcdab3f00b642375e08e9705bd726039"
+  integrity sha512-S43gOH0wA3qyv7Lc8MsQtoO/eFQNnRXi5TJmwhR4HsBUXzaJKzzwF8kOI9G22V2JBOZbsOkOgJ2Rfl9DsT2skg==
+  dependencies:
+    "@tensorflow/tfjs" "~0.14.1"
+    adm-zip "^0.4.11"
+    bindings "~1.3.0"
+    https-proxy-agent "^2.2.1"
+    node-fetch "^2.3.0"
+    progress "^2.0.0"
+    rimraf "^2.6.2"
+    tar "^4.4.6"
 
 "@tensorflow/[email protected]":
   version "0.2.1"
@@ -735,15 +748,15 @@
     rimraf "^2.6.2"
     tar "^4.4.6"
 
-"@tensorflow/tfjs@~0.14.1":
-  version "0.14.1"
-  resolved "https://registry.yarnpkg.com/@tensorflow/tfjs/-/tfjs-0.14.1.tgz#d9b198c26eb1f154c5c27228800c9df7630d2cbf"
-  integrity sha512-fo68B4FNh/JBYAjArlLFsx2etE1kxCHG11bzsy01b2jZLrM9HE+jgmiNUGuHqN7aLu/CVqB76wDAtlBT4AnPsg==
+"@tensorflow/tfjs@0.14.2", "@tensorflow/tfjs@~0.14.1":
+  version "0.14.2"
+  resolved "https://registry.yarnpkg.com/@tensorflow/tfjs/-/tfjs-0.14.2.tgz#f38fa572286dadfe981c219f5639defd586c20c4"
+  integrity sha512-d+kBdhn3L/BOIwwc44V1lUrs0O5s49ujhYXVHT9Hs6y3yq+OqPK10am16H1fNcxeMn12/3gGphebglObTD0/Sg==
   dependencies:
-    "@tensorflow/tfjs-converter" "0.7.1"
-    "@tensorflow/tfjs-core" "0.14.2"
-    "@tensorflow/tfjs-data" "0.1.4"
-    "@tensorflow/tfjs-layers" "0.9.1"
+    "@tensorflow/tfjs-converter" "0.7.2"
+    "@tensorflow/tfjs-core" "0.14.5"
+    "@tensorflow/tfjs-data" "0.1.7"
+    "@tensorflow/tfjs-layers" "0.9.2"
 
 "@types/long@^4.0.0":
   version "4.0.0"
@@ -5420,11 +5433,6 @@ user-home@^2.0.0:
   dependencies:
     os-homedir "^1.0.0"
 
-utf8@~2.1.2:
-  version "2.1.2"
-  resolved "https://registry.yarnpkg.com/utf8/-/utf8-2.1.2.tgz#1fa0d9270e9be850d9b05027f63519bf46457d96"
-  integrity sha1-H6DZJw6b6FDZsFAn9jUZv0ZFfZY=
-
 util-deprecate@^1.0.1, util-deprecate@~1.0.1:
   version "1.0.2"
   resolved "https://registry.yarnpkg.com/util-deprecate/-/util-deprecate-1.0.2.tgz#450d4dc9fa70de732762fbd2d4a28981419a0ccf"