14_8：メモリ管理

dispose

tf.Tensorやtf.Variableからdispose()を呼び出すと、そのオブジェクトがパージ(一掃)され、オブジェクトが占めていたGPUメモリが解放されます。

const x = tf.tensor2d([[1, 2], [3, 4]]);
const x_squared = x.square();
x.dispose();
x_squared.dispose();

tf.memory

この関数は、次のプロパティを持つ、現時点でのメモリ情報オブジェクトを返します。

– numBytes: 確保されたバイト数(破棄されない)
– numTensors: 確保された一意のtf.Tensorの数
– numDataBuffers: 確保された一意のデータバッファの数(破棄されない)。たとえば、myTensor.reshape(newShape)から作成されたtf.Tensorは、同じデータバッファをmyTensorと共有するので、numDataBuffersはnumTensorsよりも小さくなる。

const myTensor = tf.tensor2d([[1, 2], [3, 4], [5, 6]], [3, 2]);
const reshapedTensor = myTensor.reshape([2, 3]);

console.log('numBytes: ' + tf.memory().numBytes);
console.log('numTensors: ' + tf.memory().numTensors);
console.log('numDataBuffers: ' + tf.memory().numDataBuffers);

結果：

numBytes: 24
numTensors: 2
numDataBuffers: 1

tf.tidy

– 与えられた関数fを実行する。
– 実行後、中間的なtf.Tensorを一掃し、fによって確保されたGPUメモリを解放する(fが返す値はのぞく)。

*メモ：
– fは同期であるべきで、またPromiseを返すべきではない。
– UIを更新するコードやリモートへリクエストを行うコードはtf.tidy()の外に置くようにする。
– tf.tidy()は変数はクリーンアップしないので、手動でdispose()を呼び出す。

次のコードのa、b、oneはtidyの後、クリーンアップされます。

const result = tf.tidy(() => {
    // 4^1/2 + 1
    const x = tf.scalar(4);
    const x_sqrt = x.sqrt();
    const one = tf.scalar(1);

    console.log('numTensors (tidy): ' + tf.memory().numTensors);

    return x_sqrt.add(one);
});

console.log('numTensors (outside tidy): ' + tf.memory().numTensors);
result.print();

結果：

numTensors (tidy): 3
numTensors (outside tidy): 1
Tensor
    3

tf.keep

この関数は、tf.tidy()内で生成されたtf.Tensorを自動的に破棄されないようにします。

let x_sqrt;
const result = tf.tidy(() => {
    const x = tf.scalar(4);
    x_sqrt = tf.keep(x.sqrt());
    const one = tf.scalar(1);

    console.log('numTensors (tidy): ' + tf.memory().numTensors);

    return x_sqrt.add(one);
});

console.log('numTensors (outside tidy): ' + tf.memory().numTensors);
console.log('result: ');
result.print();
console.log('x_sqrt: ');
x_sqrt.print();

結果：

numTensors (tidy): 3
numTensors (outside tidy): 2
result: 
Tensor
    3
x_sqrt: 
Tensor
    2