How to use the @tensorflow/tfjs-node-gpu.node function in @tensorflow/tfjs-node-gpu

To help you get started, we’ve selected a few @tensorflow/tfjs-node-gpu examples, based on popular ways it is used in public projects.

Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.

github tensorflow / tfjs-examples / mnist-acgan / gan.js View on Github external
const optimizer = tf.train.adam(args.learningRate, args.adamBeta1);
  const combined = buildCombinedModel(
      args.latentSize, generator, discriminator, optimizer);

  await data.loadData();
  let {images: xTrain, labels: yTrain} = data.getTrainData();
  yTrain = tf.expandDims(yTrain.argMax(-1), -1);

  // Save the generator model once before starting the training.
  await generator.save(saveURL);

  let numTensors;
  let logWriter;
  if (args.logDir) {
    console.log(`Logging to tensorboard at logdir: ${args.logDir}`);
    logWriter = tf.node.summaryFileWriter(args.logDir);
  }

  let step = 0;
  for (let epoch = 0; epoch < args.epochs; ++epoch) {
    // Write some metadata to disk at the beginning of every epoch.
    fs.writeFileSync(
        metadataPath,
        JSON.stringify(makeMetadata(args.epochs, epoch, false)));

    const tBatchBegin = tf.util.now();

    const numBatches = Math.ceil(xTrain.shape[0] / args.batchSize);

    for (let batch = 0; batch < numBatches; ++batch) {
      const actualBatchSize = (batch + 1) * args.batchSize >= xTrain.shape[0] ?
          (xTrain.shape[0] - batch * args.batchSize) :