Tensorflow.js 內(nèi)存管理

因?yàn)門ensorFlow.js使用了GPU來加速數(shù)學(xué)運(yùn)算，因此當(dāng)tensorflow處理張量和變量時就有必要來管理GPU內(nèi)存。在TensorFlow.js中，我們可以通過dispose 和 tf.tidy這兩種方法來管理內(nèi)存。

dispose

您可以在張量或變量上調(diào)用dispose來清除它并釋放其GPU內(nèi)存：

const x = tf.tensor2d([[0.0, 2.0], [4.0, 6.0]]);
const x_squared = x.square();

x.dispose();
x_squared.dispose();

tf.tidy

進(jìn)行大量的張量操作時使用dispose可能會很麻煩。 TensorFlow.js提供了另一個函數(shù)tf.tidy，它對JavaScript中的常規(guī)范圍起到類似的作用，不同的是它針對GPU支持的張量。

tf.tidy執(zhí)行一個函數(shù)并清除所有創(chuàng)建的中間張量，釋放它們的GPU內(nèi)存。 它不清除內(nèi)部函數(shù)的返回值。

const average = tf.tidy(() => {
  const y = tf.tensor1d([1.0, 2.0, 3.0, 4.0]);
  const z = tf.ones([4]);

  return y.sub(z).square().mean();
});

average.print()

使用tf.tidy將有助于防止應(yīng)用程序中的內(nèi)存泄漏。它也可以用來更謹(jǐn)慎地控制內(nèi)存何時回收。

兩個重要的注意事項(xiàng)

傳遞給tf.tidy的函數(shù)應(yīng)該是同步的，并且不會返回Promise。我們建議在tf.tidy內(nèi)不要有更新UI或在發(fā)出遠(yuǎn)程請求的代碼。

tf.tidy不會清理變量。變量通常持續(xù)到機(jī)器學(xué)習(xí)模型的整個生命周期，因此TensorFlow.js不會清理它們，即使它們是在tidy中創(chuàng)建的。不過，您可以手動調(diào)用dispose處理它們。