TensorFlow.js optimize.minimize变量关联报错与多项式回归内存泄漏问题

Hey there! Let's tackle your two issues with adapting Daniel Shiffman's linear regression example to polynomial regression. I'll break down why each problem happens and show you a clean fix that avoids both the connection error and memory leaks.

First, Why the Non-Working Version Fails #

Your first predict function uses ys.assign(sum) to accumulate terms, but TensorFlow.js's automatic differentiation relies on building a computation graph of tensor operations—not side-effect operations like assign. When you call optimizer.minimize(), it can't trace a path from your coefficients variables to the final prediction tensor because the assign operation breaks the dependency chain. That's exactly why you get the "Cannot find a connection between any variable and the result of the loss function" error.

Second, Why the Working Version Leaks Memory #

Your second predict function fixes the dependency issue by using tensor operations instead of assign, but it creates a new tf.Variable every time it runs with let ys = tf.variable(tf.zerosLike(xs)). Unlike regular tensors, tf.Variable instances aren't automatically cleaned up by tf.tidy()—you have to manually dispose of them. Since predict gets called every frame, you're creating a new variable each time that never gets removed, leading to a memory leak.

The Fix: Pure Tensor Operations (No Extra Variables) #

The solution is to rewrite predict to use only regular tensors (no internal tf.Variable instances) and build the computation graph with standard tensor arithmetic. This keeps the dependency chain intact for the optimizer and avoids creating unmanaged variables.

Here's the corrected predict function:

function predict(x) {
  const xs = tf.tensor1d(x);
  // Initialize as a regular tensor, not a variable
  let ys = tf.zerosLike(xs);
  
  for (let i = 0; i < degree; i++) {
    const coef = coefficients[i];
    // Calculate the term: coefficient * x^(degree - i)
    const term = coef.mul(xs.pow(degree - i));
    // Accumulate using tensor addition (no assign!)
    ys = ys.add(term);
  }
  
  return ys;
}

Key Changes: #

• No internal tf.Variable: We start with tf.zerosLike(xs) (a regular tensor) instead of a variable, so there's nothing left unmanaged.
• Tensor arithmetic instead of assign: Using ys = ys.add(term) builds the computation graph correctly, so the optimizer can trace dependencies from coefficients to the final prediction.
• Clean computation graph: Every operation is a tensor transformation, which TensorFlow.js can track for automatic differentiation.

Full Corrected Code #

Here's the complete working code with the fix applied, plus some minor cleanup to ensure proper tensor disposal:

const WIDTH = 800, HEIGHT = 400;
const x_vals = [];
const y_vals = [];
let coefficients = [];
let degree = 5;
let lr = 0.2;
let optimizer = tf.train.adamax(lr);

function setup() {
  createCanvas(WIDTH, HEIGHT);
  background(0);
  initCoeffs();
  
  let up = false;
  for (let i = 0; i < WIDTH; i += WIDTH / 10) {
    x_vals.push(map(i, 0, WIDTH, -1, 1));
    y_vals.push(map((up) ? 0 : HEIGHT, 0, HEIGHT, -1, 1));
    up = !up;
  }
}

function initCoeffs() {
  for (let i = 0; i < degree; i++) {
    coefficients.push(tf.variable(tf.scalar(random(1))));
  }
}

function loss(pred, labels) {
  return tf.losses.meanSquaredError(labels, pred);
}

function predict(x) {
  const xs = tf.tensor1d(x);
  let ys = tf.zerosLike(xs);
  
  for (let i = 0; i < degree; i++) {
    const coef = coefficients[i];
    const term = coef.mul(xs.pow(degree - i));
    ys = ys.add(term);
  }
  
  return ys;
}

function draw() {
  noFill();
  background(0);
  stroke(255);
  strokeWeight(8);
  
  // Draw data points
  for (let i = 0; i < x_vals.length; i++) {
    point(
      map(x_vals[i], -1, 1, 0, WIDTH),
      map(y_vals[i], -1, 1, 0, HEIGHT)
    );
  }
  
  strokeWeight(4);
  if (x_vals.length > 0) {
    tf.tidy(() => {
      const ys = tf.tensor1d(y_vals);
      optimizer.minimize(() => loss(predict(x_vals), ys));
    });
  }
  
  // Generate and draw the regression line
  let lineX = [];
  for (let x = -1.1; x <= 1.1; x += 0.01) lineX.push(x);
  
  tf.tidy(() => {
    const lineY = predict(lineX).dataSync();
    beginShape();
    for (let i = 0; i < lineY.length; i++) {
      curveVertex(
        map(lineX[i], -1, 1, 0, WIDTH),
        map(lineY[i], -1, 1, 0, HEIGHT)
      );
    }
    endShape();
    
    // Draw line points (optional)
    stroke(200, 100, 100);
    for (let i = 0; i < lineY.length; i++) {
      point(
        map(lineX[i], -1, 1, 0, WIDTH),
        map(lineY[i], -1, 1, 0, HEIGHT)
      );
    }
  });
}

function mousePressed() {
  x_vals.push(map(mouseX, 0, WIDTH, -1, 1));
  y_vals.push(map(mouseY, 0, HEIGHT, -1, 1));
}

Additional Cleanup: #

• Wrapped the regression line drawing logic inside tf.tidy() to automatically clean up any intermediate tensors created by predict(lineX).
• Removed the manual ys.dispose() call since tf.tidy() handles it now.

HTML Dependencies (Unchanged) #

<script src="https://cdnjs.cloudflare.com/ajax/libs/p5.js/0.5.7/p5.min.js"></script>
<script src="https://cdnjs.cloudflare.com/ajax/libs/tensorflow/0.11.2/tf.min.js"></script>

This version should run without the connection error and won't leak memory—all tensors are properly managed, and the optimizer can correctly trace dependencies between your coefficients variables and the loss function.

内容的提问来源于stack exchange，提问作者Rylyn