Function main

lstm-text-generation/train_node.js:118–178

()

}

async function main() {
  const args = parseArgs();
  if (args.gpu) {
    console.log('Using GPU');
    require('@tensorflow/tfjs-node-gpu');
  } else {
    console.log('Using CPU');
    require('@tensorflow/tfjs-node');
  }

  // Create the text data object.
  let localTextDataPath = args.textDatasetPath;
  if (args.textDatasetName) {
    const textDataURL = TEXT_DATA_URLS[args.textDatasetName].url;
    localTextDataPath = path.join(os.tmpdir(), path.basename(textDataURL));
    await maybeDownload(textDataURL, localTextDataPath);
  }
  const text = fs.readFileSync(localTextDataPath, {encoding: 'utf-8'});
  const textData =
      new TextData('text-data', text, args.sampleLen, args.sampleStep);

  // Convert lstmLayerSize from string to number array before handing it
  // to `createModel()`.
  const lstmLayerSize = args.lstmLayerSize.indexOf(',') === -1 ?
      Number.parseInt(args.lstmLayerSize) :
      args.lstmLayerSize.split(',').map(x => Number.parseInt(x));

  const model = createModel(
      textData.sampleLen(), textData.charSetSize(), lstmLayerSize);
  compileModel(model, args.learningRate);

  // Get a seed text for display in the course of model training.
  const [seed, seedIndices] = textData.getRandomSlice();
  console.log(`Seed text:\n"${seed}"\n`);

  const DISPLAY_TEMPERATURES = [0, 0.25, 0.5, 0.75];

  let epochCount = 0;
  await fitModel(
      model, textData, args.epochs, args.examplesPerEpoch, args.batchSize,
      args.validationSplit, {
        onTrainBegin: async () => {
          epochCount++;
          console.log(`Epoch ${epochCount} of ${args.epochs}:`);
        },
        onTrainEnd: async () => {
          DISPLAY_TEMPERATURES.forEach(async temperature => {
            const generated = await generateText(
                model, textData, seedIndices, args.displayLength, temperature);
            console.log(
                `Generated text (temperature=${temperature}):\n` +
                `"${generated}"\n`);
          });
        }
      });

  if (args.savePath != null && args.savePath.length > 0) {
    await model.save(`file://${args.savePath}`);