1 |
|
2 |
|
3 |
|
4 |
|
5 |
|
6 |
|
7 |
|
8 |
|
9 |
|
10 |
|
11 |
|
12 |
|
13 |
|
14 |
|
15 |
|
16 |
|
17 |
|
18 |
|
19 | import './register_all_kernels';
|
20 |
|
21 | import * as tf from '@tensorflow/tfjs';
|
22 | import * as path from 'path';
|
23 |
|
24 | import {ProgbarLogger} from './callbacks';
|
25 | import {nodeFileSystemRouter} from './io/file_system';
|
26 | import * as nodeIo from './io/index';
|
27 | import {NodeJSKernelBackend} from './nodejs_kernel_backend';
|
28 | import {TFJSBinding} from './tfjs_binding';
|
29 | import * as nodeVersion from './version';
|
30 |
|
31 |
|
32 | const binary = require('@mapbox/node-pre-gyp');
|
33 | const bindingPath =
|
34 | binary.find(path.resolve(path.join(__dirname, '/../package.json')));
|
35 |
|
36 |
|
37 |
|
38 | const fs = require('fs');
|
39 | if (!fs.existsSync(bindingPath)) {
|
40 | throw new Error(
|
41 | `The Node.js native addon module (tfjs_binding.node) can not ` +
|
42 | `be found at path: ` + String(bindingPath) + `. \nPlease run command ` +
|
43 | `'npm rebuild @tensorflow/tfjs-node` +
|
44 | (String(bindingPath).indexOf('tfjs-node-gpu') > 0 ? `-gpu` : ``) +
|
45 | ` --build-addon-from-source' to ` +
|
46 | `rebuild the native addon module. \nIf you have problem with building ` +
|
47 | `the addon module, please check ` +
|
48 | `https://github.com/tensorflow/tfjs/blob/master/tfjs-node/` +
|
49 | `WINDOWS_TROUBLESHOOTING.md or file an issue.`);
|
50 | }
|
51 |
|
52 | const bindings = require(bindingPath);
|
53 |
|
54 |
|
55 | export const version = {
|
56 | ...tf.version,
|
57 | 'tfjs-node': nodeVersion.version
|
58 | };
|
59 | export const io = {
|
60 | ...tf.io,
|
61 | ...nodeIo
|
62 | };
|
63 |
|
64 |
|
65 | export * from '@tensorflow/tfjs';
|
66 | export * from './node';
|
67 |
|
68 |
|
69 | const pjson = require('../package.json');
|
70 |
|
71 |
|
72 | tf.registerBackend('tensorflow', () => {
|
73 | return new NodeJSKernelBackend(bindings as TFJSBinding, pjson.name);
|
74 | }, 3 );
|
75 |
|
76 | const success = tf.setBackend('tensorflow');
|
77 | if (!success) {
|
78 | throw new Error(`Could not initialize TensorFlow backend.`);
|
79 | }
|
80 |
|
81 |
|
82 | tf.io.registerLoadRouter(nodeFileSystemRouter);
|
83 | tf.io.registerSaveRouter(nodeFileSystemRouter);
|
84 |
|
85 |
|
86 | tf.registerCallbackConstructor(1, ProgbarLogger);
|