Function buildGRUModel

jena-weather/models.js:144–157

(inputShape, dropout, recurrentDropout)

 * @returns {tf.LayersModel} A TensorFlow.js GRU model.
 */
export function buildGRUModel(inputShape, dropout, recurrentDropout) {
  // TODO(cais): Recurrent dropout is currently not fully working.
  //   Make it work and add a flag to train-rnn.js.
  const model = tf.sequential();
  const rnnUnits = 32;
  model.add(tf.layers.gru({
    units: rnnUnits,
    inputShape,
    dropout: dropout || 0,
    recurrentDropout: recurrentDropout || 0
  }));
  model.add(tf.layers.dense({units: 1}));
  return model;
}

/**
 * Build a model for the temperature-prediction problem.