Function getIrisData

iris-fitDataset/data.js:106–137

(testSplit)

 *   - Labels as a rank-1 `Tensor` in one-hot format.
 */
export async function getIrisData(testSplit) {
  // TODO(bileschi): Update shuffle etc. to use the tf.data API calls once
  // it is possible to cache the results for performance and train-test split
  // stability across epochs.  Once caching is available, perform batching first
  // and then map the preprocessing functions across the batches.
  // https://github.com/tensorflow/tfjs/issues/1025

  // Shuffle a copy of the raw data.
  const shuffled = IRIS_RAW_DATA.slice();
  tf.util.shuffle(shuffled);
  // Split the data into training and testing portions.
  const numTestExamples = Math.round(IRIS_RAW_DATA.length * testSplit);
  const numTrainExamples = IRIS_RAW_DATA.length - numTestExamples;
  const train = shuffled.slice(0, numTrainExamples);
  const test = shuffled.slice(numTrainExamples);
  // Split the data into into X & y and apply feature mapping transformations
  const trainX = tf.data.array(train.map(r => r.slice(0, 4)));
  const testX = tf.data.array(test.map(r => r.slice(0, 4)));
  // TODO(we should be able to just directly use tensors built from oneHot here
  // instead of converting to tensor and back using datasync & Array.from.
  // This causes an internal disposal error however.
  // See https://github.com/tensorflow/tfjs/issues/1071
  //
  // const trainY = tf.data.array(train.map(r => tf.oneHot([r[4]], 3)));
  // const testY = tf.data.array(test.map(r => tf.oneHot([r[4]], 3)));
  const trainY = tf.data.array(train.map(r => flatOneHot(r[4])));
  const testY = tf.data.array(test.map(r => flatOneHot(r[4])));
  // Recombine the X and y portions of the data.
  const trainDataset = tf.data.zip({xs: trainX, ys: trainY});
  const testDataset = tf.data.zip({xs: testX, ys: testY});
  return [trainDataset, testDataset];
}