Allow using GPU

package.json

@@ -6,6 +6,7 @@
   "private": true,
   "dependencies": {
     "@tensorflow/tfjs-node": "^4.1.0",
+    "@tensorflow/tfjs-node-gpu": "^4.1.0",
     "tslib": "^2.4.1"
   },
   "devDependencies": {
@@ -30,4 +31,3 @@
     "typescript": "4.8.4"
   }
 }
-

packages/tfjs-node-helpers/src/classification/binary-classification-trainer.ts

@@ -22,8 +22,10 @@ import { Metric } from '../testing/metric';
 import { MetricCalculator } from '../testing/metric-calculator';
 import { binarize } from '../utils/binarize';
 import { FeatureNormalizer } from '../feature-engineering/feature-normalizer';
+import { switchHardwareUsage } from '../utils/switch-hardware-usage';
 
 export type BinaryClassificationTrainerOptions = {
+  shouldUseGPU?: boolean;
   batchSize?: number;
   epochs?: number;
   patience?: number;
@@ -53,6 +55,8 @@ export class BinaryClassificationTrainer {
   protected static DEFAULT_PATIENCE: number = 20;
 
   constructor(options: BinaryClassificationTrainerOptions) {
+    switchHardwareUsage(options.shouldUseGPU);
+
     this.batchSize = options.batchSize ?? BinaryClassificationTrainer.DEFAULT_BATCH_SIZE;
     this.epochs = options.epochs ?? BinaryClassificationTrainer.DEFAULT_EPOCHS;
     this.patience = options.patience ?? BinaryClassificationTrainer.DEFAULT_PATIENCE;

packages/tfjs-node-helpers/src/classification/binary-classifier.ts

@@ -1,13 +1,17 @@
 import { LayersModel, loadLayersModel, tensor, Tensor } from '@tensorflow/tfjs-node';
+import { switchHardwareUsage } from '../utils/switch-hardware-usage';
 
 export type BinaryClassifierOptions = {
   model: LayersModel;
+  shouldUseGPU?: boolean;
 };
 
 export class BinaryClassifier {
   protected model?: LayersModel;
 
   constructor(options?: BinaryClassifierOptions) {
+    switchHardwareUsage(options?.shouldUseGPU);
+
     this.model = options?.model;
   }
 

packages/tfjs-node-helpers/src/utils/switch-hardware-usage.ts

@@ -0,0 +1,9 @@
+export const switchHardwareUsage = (shouldUseGPU?: boolean): void => {
+  if (shouldUseGPU) {
+    console.log('Attempting to use GPU.');
+    require('@tensorflow/tfjs-node-gpu');
+  } else {
+    console.log('Using CPU.');
+    require('@tensorflow/tfjs-node');
+  }
+};