Function sampleFromMnistData

mnist-acgan/web-data.js:171–210

(numExamplesPerClass)

 *   [numExamplesPerClass * 10, 28, 28, 1].
 */
export function sampleFromMnistData(numExamplesPerClass) {
  tf.util.assert(
      numExamplesPerClass <= mnistNumExamples / 10,
      `Requested too many examples per class ` +
          `(${numExamplesPerClass} > ${mnistNumExamples / 10})`);

  tf.util.shuffle(mnistIndices);
  const indicesByClass = [];
  for (let i = 0; i < NUM_CLASSES; ++i) {
    indicesByClass.push([]);
  }

  for (let i = 0; i < mnistIndices.length; ++i) {
    if (indicesByClass[mnistLabels[mnistIndices[i]]].length >=
        numExamplesPerClass) {
      continue;
    }
    indicesByClass[mnistLabels[mnistIndices[i]]].push(mnistIndices[i]);

    let minLength = Infinity;
    indicesByClass.forEach(indicesArray => {
      if (indicesArray.length < minLength) {
        minLength = indicesArray.length;
      }
    });
    if (minLength >= numExamplesPerClass) {
      break;
    }
  }

  return tf.tidy(() => {
    let rowsToCombine = [];
    indicesByClass.forEach(classIndices => {
      const classImages = tf.gather(mnistImages, classIndices);
      const rowOfExamples = tf.concat(classImages.unstack(), 0);
      rowsToCombine.push(rowOfExamples);
    });
    return tf.concat(rowsToCombine, 1);
  });
}