Method makeTrainFunction

tfjs-layers/src/engine/training.ts:1298–1382

* Creates a function that performs the following actions: * * 1. computes the losses * 2. sums them to get the total loss * 3. call the optimizer computes the gradients of the LayersModel's * trainable weights w.r.t. the total loss and update the variables * 4. calculates the me

()

   * 5. returns the values of the losses and metrics.
   */
  protected makeTrainFunction(): (data: Tensor[]) => Scalar[] {
    return (data: Tensor[]) => {
      const lossValues: Scalar[] = [];

      const inputs = data.slice(0, this.inputs.length);
      const targets = data.slice(
          this.inputs.length, this.inputs.length + this.outputs.length);
      const sampleWeights = data.slice(
          this.inputs.length + this.outputs.length,
          this.inputs.length + this.outputs.length * 2);

      const metricsValues: Scalar[] = [];

      // Create a function that computes the total loss based on the
      // inputs. This function is used for obtaining gradients through
      // backprop.
      const totalLossFunction = () => {
        const feeds = [];
        for (let i = 0; i < this.inputs.length; ++i) {
          feeds.push({key: this.inputs[i], value: inputs[i]});
        }
        const feedDict = new FeedDict(feeds);
        const outputs =
            execute(this.outputs, feedDict, {'training': true}) as Tensor[];
        // TODO(cais): Take care of the case of multiple outputs from a
        //   single layer?

        let totalLoss: Tensor;
        for (let i = 0; i < this.lossFunctions.length; ++i) {
          const lossFunction = this.lossFunctions[i];
          let loss = lossFunction(targets[i], outputs[i]);
          if (sampleWeights[i] != null) {
            loss = computeWeightedLoss(loss, sampleWeights[i]);
          }

          // TODO(cais): push Scalar instead.
          const meanLoss: Scalar = tfc.mean(loss);
          // TODO(cais): Use a scope() instead, to avoid ownership.
          lossValues.push(meanLoss);
          if (i === 0) {
            totalLoss = loss;
          } else {
            totalLoss = tfc.add(totalLoss, loss);
          }
        }

        // Compute the metrics.
        // TODO(cais): These should probably be calculated outside
        //   totalLossFunction to benefit speed?
        for (let i = 0; i < this.metricsTensors.length; ++i) {
          let weightedMetric: Scalar;

          if (this.outputs.length > 1 && i < this.outputs.length) {
            weightedMetric = lossValues[i];
          } else {
            const metric = this.metricsTensors[i][0];
            const outputIndex = this.metricsTensors[i][1];
            weightedMetric =