Function run

date-conversion-attention/train.js:177–240

()

}

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

  const model = createModel(
      dateFormat.INPUT_VOCAB.length, dateFormat.OUTPUT_VOCAB.length,
      dateFormat.INPUT_LENGTH, dateFormat.OUTPUT_LENGTH);
  model.summary();

  const {
    trainEncoderInput,
    trainDecoderInput,
    trainDecoderOutput,
    valEncoderInput,
    valDecoderInput,
    valDecoderOutput,
    testDateTuples
  } = generateDataForTraining(args.trainSplit, args.valSplit);

  await model.fit(
      [trainEncoderInput, trainDecoderInput], trainDecoderOutput, {
        epochs: args.epochs,
        batchSize: args.batchSize,
        shuffle: true,
        validationData: [[valEncoderInput, valDecoderInput], valDecoderOutput],
        callbacks: args.logDir == null ? null :
            tfn.node.tensorBoard(args.logDir, {updateFreq: args.logUpdateFreq})
      });

  // Save the model.
  if (args.savePath != null && args.savePath.length) {
    if (!fs.existsSync(args.savePath)) {
      shelljs.mkdir('-p', args.savePath);
    }
    const saveURL = `file://${args.savePath}`
    await model.save(saveURL);
    console.log(`Saved model to ${saveURL}`);
  }

  // Run seq2seq inference tests and print the results to console.
  const numTests = 10;
  for (let n = 0; n < numTests; ++n) {
    for (const testInputFn of dateFormat.INPUT_FNS) {
      const inputStr = testInputFn(testDateTuples[n]);
      console.log('\n-----------------------');
      console.log(`Input string: ${inputStr}`);
      const correctAnswer =
          dateFormat.dateTupleToYYYYDashMMDashDD(testDateTuples[n]);
      console.log(`Correct answer: ${correctAnswer}`);

      const {outputStr} = await runSeq2SeqInference(model, inputStr);