Skip to content

Commit

Permalink
Support a GPU build (#184)
Browse files Browse the repository at this point in the history
  • Loading branch information
Kevin Scott committed Jan 29, 2022
1 parent 1ee6f9f commit 79374d3
Show file tree
Hide file tree
Showing 7 changed files with 2,419 additions and 2,407 deletions.
1 change: 1 addition & 0 deletions examples/nodejs/src/index.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ const tf = require('@tensorflow/tfjs-node');

const app = express();


app.get("/", async (req, res) => {
const upscaledImage = await getUpscaledImage();
res.set('Content-Type', 'image/png');
Expand Down
8 changes: 3 additions & 5 deletions package.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@
]
},
"scripts": {
"example": "node ./scripts/start_example.js",
"example": "node ./scripts/start-example.js",
"update:version": "node ./scripts/version.js",
"format": "yarn lerna run format",
"lint": "yarn lerna run lint",
Expand All @@ -38,10 +38,6 @@
"postversion": "yarn build",
"bootstrap": "yarn lerna bootstrap --hoist",
"build": "yarn lerna run build",
"build:cjs": "yarn lerna run build:cjs",
"build:umd": "yarn lerna run build:umd",
"build:esm": "yarn lerna run build:esm",
"watch:esm": "yarn lerna run watch:esm",
"server": "yarn serve -p 8000 packages/upscalerjs/test/lib/webpack-bundler/dist/",
"set:tfjs": "yarn lerna run set:tfjs",
"test": "yarn test:unit",
Expand Down Expand Up @@ -92,6 +88,7 @@
"rollup-plugin-node-resolve": "^5.2.0",
"selenium-webdriver": "^4.0.0-beta.1",
"serve": "^11.3.2",
"tree-kill": "^1.2.2",
"ts-jest": "^26.1.4",
"ts-node": "^10.4.0",
"tslint": "^6.1.2",
Expand All @@ -102,6 +99,7 @@
"yargs": "^16.2.0"
},
"dependencies": {
"chokidar": "^3.5.3",
"tensor-as-base64": "^0.1.1"
},
"packageManager": "yarn@3.1.0"
Expand Down
28 changes: 15 additions & 13 deletions packages/upscalerjs/package.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,14 +3,15 @@
"version": "0.10.0-canary.1",
"description": "Increase Image resolution with Tensorflow.js",
"exports": {
"./node": "./dist/cjs/cjs.js",
".": "./dist/esm/index.js"
"./node": "./dist/node/cjs/cjs.js",
"./node-gpu": "./dist/node-gpu/cjs/cjs.js",
".": "./dist/browser/index.js"
},
"engines": {
"node": ">=0.13"
},
"module": "dist/esm/index.js",
"types": "dist/esm/index.d.ts",
"module": "dist/browser/index.js",
"types": "dist/browser/index.d.ts",
"keywords": [
"super resolution",
"image",
Expand All @@ -32,9 +33,9 @@
"scripts": {
"build": "rimraf dist && yarn build:node && yarn build:browser",

"build:browser": "yarn set:tfjs '@tensorflow/tfjs' && yarn lerna run build:umd && yarn lerna run build:esm",
"build:node": "yarn set:tfjs '@tensorflow/tfjs-node' && yarn lerna run build:cjs",
"build:node-gpu": "yarn set:tfjs '@tensorflow/tfjs-node-gpu' && yarn lerna run build:cjs",
"build:browser": "rimraf dist/browser && mkdir -p dist/browser && yarn set:tfjs '@tensorflow/tfjs' && yarn lerna run build:umd && yarn lerna run build:esm && mv dist/esm dist/browser/esm",
"build:node": "rimraf dist/node && mkdir -p dist/node && yarn set:tfjs '@tensorflow/tfjs-node' && yarn lerna run build:cjs && mv dist/cjs dist/node/cjs",
"build:node-gpu": "rimraf dist/node-gpu && mkdir -p dist/node-gpu && yarn set:tfjs '@tensorflow/tfjs-node-gpu' && yarn lerna run build:cjs && mv dist/cjs dist/node-gpu/cjs",

"build:cjs": "tsc --module commonjs --target es5 --outDir dist/cjs",
"build:esm": "tsc --module esnext --moduleResolution node --target esnext --outDir dist/esm",
Expand All @@ -54,12 +55,13 @@
"lint": "echo 'Linting code' && tslint -p tsconfig.json",
"test": "jest --config jestconfig.json",

"watch": "watch 'yarn build' --filter=../../scripts/watch-filter.js",
"watch:esm": "watch 'yarn build:esm' --filter=../../scripts/watch-filter.js",
"watch:umd": "watch 'yarn build:umd' --filter=../../scripts/watch-filter.js",
"watch:cjs": "watch 'yarn build:cjs' --filter=../../scripts/watch-filter.js",
"watch:browser": "watch 'yarn build:browser' --filter=../../scripts/watch-filter.js",
"watch:node": "watch 'yarn build:node' --filter=../../scripts/watch-filter.js"
"watch": "chokidar \"src\" -c \"yarn build\" -i \"test.ts\" -i \"generated.ts\"",
"watch:esm": "chokidar \"src\" -c \"yarn build:esm\" -i \"test.ts\" -i \"generated.ts\"",
"watch:umd": "chokidar \"src\" -c \"yarn build:umd\" -i \"test.ts\" -i \"generated.ts\"",
"watch:cjs": "chokidar \"src\" -c \"yarn build:cjs\" -i \"test.ts\" -i \"generated.ts\"",
"watch:browser": "chokidar \"src\" -c \"yarn build:browser\" -i \"test.ts\" -i \"generated.ts\"",
"watch:node": "chokidar \"src\" -c \"yarn build:node\" -i \"test.ts\" -i \"generated.ts\"",
"watch:node-gpu": "chokidar \"src\" -c \"yarn build:node-gpu\" -i \"test.ts\" -i \"generated.ts\""
},
"repository": {
"type": "git",
Expand Down
2 changes: 1 addition & 1 deletion packages/upscalerjs/src/models/idealo-gans.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,7 @@ class PixelShuffle extends tf.layers.Layer {
const config: IIntermediaryModelDefinition = {
urlPath: 'idealo/gans',
scale: 4,
preprocess: (image) => tf.mul(image, 1 / 255) as tf.Tensor4D,
preprocess: (image) => tf.mul(image, 1 / 255),
postprocess: (output: tf.Tensor3D) => tf.mul(output.clipByValue(0, 1), 255),
customLayers: [MultiplyBeta, PixelShuffle],
};
Expand Down
39 changes: 31 additions & 8 deletions scripts/start_example.js → scripts/start-example.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,11 @@
*/

const fs = require('fs');
const path = require('path');
const { spawn } = require('child_process');
const kill = require('tree-kill');
const chokidar = require('chokidar');


const exampleDirectory = process.argv.pop();

Expand All @@ -14,20 +18,24 @@ try {
process.exit(1)
}

const runProcess = (command, args, name) => {
const spawnedProcess = spawn(command, args);
const runProcess = (command, args = []) => {
const spawnedProcess = spawn(command, args, { shell: true });

spawnedProcess.stdout.on('data', (data) => {
process.stdout.write(`${name}: ${data.toString()}`);
process.stdout.write(data.toString());
});

spawnedProcess.stderr.on('data', (data) => {
process.stderr.write(`${name}: ${data.toString()}`);
process.stderr.write(data.toString());
});

spawnedProcess.on('exit', (data) => {
process.stderr.write(`${name}: ${data.toString()}`);
if (data) {
process.stderr.write(data.toString());
}
});

return spawnedProcess;
}

const getPlatform = () => {
Expand All @@ -37,7 +45,7 @@ const getPlatform = () => {
} else if (deps.includes('@tensorflow/tfjs-node')) {
return 'node';
} else if (deps.includes('@tensorflow/tfjs-node-gpu')) {
return 'node';
return 'node-gpu';
}

throw new Error('Could not determine valid TFJS dependency in example package.json')
Expand All @@ -46,11 +54,26 @@ const getPlatform = () => {
// get package name from directory
const packageJSON = JSON.parse(fs.readFileSync(`./examples/${exampleDirectory}/package.json`, 'utf8'));
const exampleName = packageJSON.name;
const examplePath = path.resolve(__dirname, `../examples/${exampleName}`);
const platform = getPlatform();

const main = () => {
runProcess('yarn', ['workspace', exampleName, 'start'], 'example');
runProcess('yarn', ['workspace', 'upscaler', `watch:${platform}`], 'upscaler');
let process;
const onChange = () => {
if (process) {
kill(process.pid);
}
const args = ['yarn', 'workspace', 'upscaler', `build:${platform}`, '&&', 'yarn', 'workspace', exampleName, 'start'];
process = runProcess(args.join(' '));
};
chokidar.watch(path.resolve(__dirname, '../packages/upscalerjs/src'), {
ignored: /((^|[\/\\])\..|test.ts|generated.ts)/, // ignore dotfiles
persistent: true
}).on('all', onChange);
chokidar.watch(examplePath, {
ignored: /(^|[\/\\])\../, // ignore dotfiles
persistent: true
}).on('all', onChange);
};

main();
25 changes: 0 additions & 25 deletions scripts/watch-filter.js
View file

This file was deleted.

Loading

0 comments on commit 79374d3

Please sign in to comment.