-
Notifications
You must be signed in to change notification settings - Fork 208
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
feat: add tvm wasm export #582
Changes from all commits
a32c280
cda9a77
aa6eb93
f9a2b71
824ae18
eace176
ff80ea1
7627d10
36cb1ae
1cd0241
0c968f6
b07993c
1265a49
9696f42
f11ae7a
b79d1cd
14ed9ee
2d12f51
9330636
28a0142
6d54394
ea686c7
cb81a8e
7139b23
cc004ea
def08c4
187d110
1edf04e
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -175,7 +175,7 @@ export class PipelineService { | |
const verifyPlugin = (name: string): void => { | ||
if (!plugins[name]) { | ||
this.runnableMap[job.id].destroy(); | ||
throw new TypeError(`"${name}" plugin is required`); | ||
throw new TypeError(`'${name}' plugin is required`); | ||
} | ||
}; | ||
const dispatchJobEvent = (jobStatus: PipelineStatus, step?: PluginTypeI, stepAction?: 'start' | 'end') => { | ||
|
@@ -271,7 +271,7 @@ export class PipelineService { | |
datasetProcess, | ||
pipeline, | ||
workingDir: runnable.workingDir, | ||
template: 'node' // set node by default | ||
template: process.env.WASM ? 'wasm' : 'node' // set node by default | ||
}); | ||
|
||
await JobModel.saveJob(job); | ||
|
@@ -317,6 +317,66 @@ export class PipelineService { | |
return path.join(CoreConstants.PIPCOOK_RUN, id, 'output.tar.gz'); | ||
} | ||
|
||
// private _generateWASMOutput(dist: string, opts: GenerateOptions, fileQueue: Array<Promise<void | string>>): void { | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. @WenheLI @FeelyChau Boa is working with |
||
// const relay = boa.import('tvm.relay'); | ||
// const emcc = boa.import('tvm.contrib.emcc'); | ||
// const keras = boa.import('tensorflow.keras'); | ||
// const {dict, open} = boa.builtins(); | ||
|
||
// // download tvm runtime from oss | ||
// const tvmjsPromise = execAsync(`wget http://ai-sample.oss-cn-hangzhou.aliyuncs.com/tvmjs/dist/tvmjs.bundle.js`, {cwd: dist}); | ||
// fileQueue.push(tvmjsPromise); | ||
|
||
// const model = keras.models.load_model(path.join(opts.modelPath, 'model.h5')); | ||
|
||
// const inputName = 'input_1'; | ||
// const inputShape = model.layers[0].input_shape[0]; | ||
// const shape = [1]; | ||
// shape.push(inputShape[3]); | ||
// shape.push(inputShape[1]); | ||
// shape.push(inputShape[2]); | ||
|
||
// const [ mod, params ] = relay.frontend.from_keras(model, dict(boa.kwargs({[inputName]: shape}))); | ||
// const [ graph, lib, param ] = relay.build(mod, boa.kwargs({ | ||
// params, | ||
// target: 'llvm -mtriple=wasm32--unknown-emcc -system-lib' | ||
// })); | ||
|
||
// lib.save(path.join(dist, 'model.bc')); | ||
|
||
// const jsonWriter = open(path.join(dist, 'modelDesc.json'), 'w'); | ||
// jsonWriter.write(graph); | ||
// const paramWriter = open(path.join(dist, 'modelParams.parmas'), 'wb'); | ||
// paramWriter.write(relay.save_param_dict(param)); | ||
// emcc.create_tvmjs_wasm(path.join(dist, 'model.wasi.js'), path.join(dist, 'model.bc'), boa.kwargs({ | ||
// options: ['-O3', '-std=c++14', '-Wno-ignored-attributes', '-s', 'ALLOW_MEMORY_GROWTH=1', '-s', 'STANDALONE_WASM=1', '-s', 'ERROR_ON_UNDEFINED_SYMBOLS=0', '-s', 'ASSERTIONS=1', '--no-entry', '--pre-js', './packages/daemon/binary/preload.js'] | ||
// })); | ||
|
||
// const templateHead = `function EmccWASI() {`; | ||
// const templateTail = ` | ||
// this.Module = Module; | ||
// this.start = Module.wasmLibraryProvider.start; | ||
// this.imports = Module.wasmLibraryProvider.imports; | ||
// this.wasiImport = this.imports['wasi_snapshot_preview1']; | ||
// } | ||
|
||
// if (typeof module !== 'undefined' && module.exports) { | ||
// module.exports = EmccWASI; | ||
// } | ||
// `; | ||
|
||
// const result = templateHead + open(path.join(dist, 'model.wasi.js')).read() + templateTail; | ||
// const resultWriter = open(path.join(dist, 'model.wasi.js'), 'w'); | ||
// resultWriter.write(result); | ||
|
||
// const jsonPromise = fs.writeJSON(path.join(dist, 'modelSpec.json'), { | ||
// shape, | ||
// inputName | ||
// }); | ||
|
||
// fileQueue.push(jsonPromise); | ||
// } | ||
|
||
/** | ||
* Generate the output package for a given job. | ||
* @param job the job model for output. | ||
|
@@ -327,18 +387,34 @@ export class PipelineService { | |
const dist = path.join(opts.workingDir, 'output'); | ||
await fs.remove(dist); | ||
await fs.ensureDir(dist); | ||
|
||
const fileQueue: Array<Promise<void | string>> = new Array(); | ||
|
||
// Only support tensorflow at this moment. | ||
// if (opts.template === 'wasm' && opts.modelPlugin.name.includes('tensorflow')) { | ||
// this._generateWASMOutput(dist, opts, fileQueue); | ||
// } | ||
|
||
await execAsync('npm init -y', { cwd: dist }); | ||
|
||
// post processing the package.json | ||
const projPackage = await fs.readJSON(dist + '/package.json'); | ||
projPackage.dependencies = { | ||
[opts.modelPlugin.name]: opts.modelPlugin.version, | ||
}; | ||
projPackage.scripts = { | ||
postinstall: 'node boapkg.js' | ||
}; | ||
if (opts.dataProcess) { | ||
projPackage.dependencies[opts.dataProcess.name] = opts.dataProcess.version; | ||
|
||
if (opts.template === 'node') { | ||
projPackage.dependencies = { | ||
[opts.modelPlugin.name]: opts.modelPlugin.version, | ||
}; | ||
projPackage.scripts = { | ||
postinstall: 'node boapkg.js' | ||
}; | ||
if (opts.dataProcess) { | ||
projPackage.dependencies[opts.dataProcess.name] = opts.dataProcess.version; | ||
} | ||
} else { | ||
projPackage.main = 'index.js'; | ||
projPackage.dependencies = { | ||
ws: '^7.3.1' | ||
}; | ||
} | ||
|
||
const jsonWriteOpts = { spaces: 2 } as fs.WriteOptions; | ||
|
@@ -347,18 +423,21 @@ export class PipelineService { | |
output: job, | ||
}; | ||
|
||
await Promise.all([ | ||
if (opts.template === 'node') { | ||
// copy base components | ||
fs.copy(opts.modelPath, dist + '/model'), | ||
fs.copy(path.join(__dirname, `../../templates/${opts.template}/predict.js`), `${dist}/index.js`), | ||
fs.copy(path.join(__dirname, '../../templates/boapkg.js'), `${dist}/boapkg.js`), | ||
// copy logs | ||
fs.copy(opts.workingDir + '/logs', `${dist}/logs`), | ||
// write package.json | ||
fs.outputJSON(dist + '/package.json', projPackage, jsonWriteOpts), | ||
// write metadata.json | ||
fs.outputJSON(dist + '/metadata.json', metadata, jsonWriteOpts), | ||
]); | ||
fileQueue.push(fs.copy(opts.modelPath, dist + '/model')); | ||
fileQueue.push(fs.copy(path.join(__dirname, '../../templates/boapkg.js'), `${dist}/boapkg.js`)); | ||
} | ||
|
||
fileQueue.push(fs.copy(path.join(__dirname, `../../templates/${opts.template}/predict.js`), `${dist}/index.js`)); | ||
// copy logs | ||
fileQueue.push(fs.copy(opts.workingDir + '/logs', `${dist}/logs`)); | ||
// write package.json | ||
fileQueue.push(fs.outputJSON(dist + '/package.json', projPackage, jsonWriteOpts)); | ||
// write metadata.json | ||
fileQueue.push(fs.outputJSON(dist + '/metadata.json', metadata, jsonWriteOpts)); | ||
|
||
await Promise.all(fileQueue); | ||
console.info(`trained the model to ${dist}`); | ||
|
||
// packing the output directory. | ||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,41 @@ | ||
const tvmjs = require("./tvmjs.bundle"); | ||
const EmccWASI = require("./model.wasi"); | ||
const fs = require('fs'); | ||
const modelSpec = require("./modelSpec.json"); | ||
|
||
const loadModel = async () => { | ||
const wasmSource = fs.readFileSync('./model.wasi.wasm'); | ||
const tvm = await tvmjs.instantiate(wasmSource, new EmccWASI()); | ||
|
||
const graph = JSON.parse(fs.readFileSync('./modelDesc.json')); | ||
const param = new Uint8Array(fs.readFileSync('./modelParams.parmas')); | ||
|
||
const ctx = tvm.cpu(0); | ||
const sysLib = tvm.systemLib(); | ||
model = tvm.createGraphRuntime(JSON.stringify(graph), sysLib, ctx); | ||
model.loadParams(param); | ||
|
||
return {model, tvm, ctx}; | ||
} | ||
|
||
let model, tvm, ctx; | ||
|
||
const predict = async (input) => { | ||
if (!model) { | ||
const rets = await loadModel(); | ||
model = rets.model; | ||
tvm = rets.tvm; | ||
ctx = rets.ctx; | ||
} | ||
|
||
const inputData = tvm.empty(modelSpec.shape, "float32", tvm.cpu()); | ||
const output = model.getOutput(0); | ||
inputData.copyFrom(input); | ||
model.setInput(modelSpec.inputName, inputData); | ||
model.run(); | ||
await ctx.sync(); | ||
console.log(output.toArray()) | ||
return output.toArray(); | ||
} | ||
|
||
module.exports = predict; | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. EOF |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Shall we include the emcc/emsdk inside Pipcook?