Утечка памяти в Tensorflow.js: как убрать неиспользуемые тензоры? - PullRequest
Новая
       19

Утечка памяти в Tensorflow.js: как убрать неиспользуемые тензоры?

2 голосов
/ 09 июня 2019

Я пишу скрипт, который иногда пропускает тензоры.Это может происходить во многих случаях, например, когда я тренирую нейронную сеть, но тренировка падает.В этом случае обучение прерывается и не удастся правильно расположить тензоры.Это приводит к утечке памяти, которую я пытаюсь устранить, удаляя неиспользуемые тензоры.

Пример

В приведенном ниже фрагменте я тренируюсь два (очень простые) модели.Первый запуск будет работать и приведет к отсутствию утечек тензоров (количество тензоров до тренировки = количество тензоров после тренировки).Во второй раз я использую недопустимый слой reshape, чтобы вызвать сбой во время тренировки.Поэтому выдается ошибка, и тензоры из набора данных (я полагаю?) Будут расположены неправильно.Код является примером, показывающим, как могут вытекать тензоры.

async function train(shouldCrash) {
  console.log(`Training, shouldCrash=${shouldCrash}`);
  const dataset = tf.data.zip({ // setup data
    xs: tf.data.array([[1],[1]]),
    ys: tf.data.array([1]),
  }).batch(1);

  const model = tf.sequential({ // setup model
    layers: [
      tf.layers.dense({units: 1, inputShape: [1]}),
      tf.layers.reshape({targetShape: [(shouldCrash ? 2 : 1)]}), // use invalid shape when crashing
    ],
  });
  model.compile({ optimizer: 'sgd', loss: 'meanSquaredError' });
  console.log('  Tensors before:', tf.memory().numTensors);
  try {
    const history = await model.fitDataset(dataset, { epochs: 1 });
  } catch (err) {
    console.log(`    Error: ${err.message}`);
  }
  console.log('  Tensors after:', tf.memory().numTensors);
}

(async () => {
  await train(false); // normal training
  await train(true); // training with error
})();
<script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@1.1.2/dist/tf.min.js"></script>

Вопрос

Существует tf.tidy, который помогает мне в некоторых случаяхизбавиться от неиспользуемых тензоров, но его можно использовать только для синхронных вызовов функций.Поэтому его нельзя использовать при вызове await model.fitDataset(...).

. Есть ли способ избавиться от неиспользуемых тензоров?В качестве альтернативы, есть ли способ удалить все существующие тензоры на странице (без перезагрузки)?

1 Ответ

2 голосов
/ 10 июня 2019

Согласно документации, функция, предоставленная для tf.tidy "не должна возвращать обещание". Внутренне, tf backend располагает всеми тензорами, используемыми при подгонке модели. Вот почему tf.fit не следует размещать внутри tf.tidy. Чтобы утилизировать разбившуюся модель, можно позвонить на модель tf.dispose.

Это правда, что в настоящее время, похоже, происходит утечка памяти, но сбой модели во время определения модели является плохой реализацией. Это не должно происходить в правильном сценарии, так как можно проверить, соответствуют ли заданные параметры тому, что должно быть входом для слоев. Например, перед построением модели можно избежать изменения формы от 2 до 1. Чтобы предотвратить утечку памяти. 

async function train(shouldCrash) {
  console.log(`Training, shouldCrash=${shouldCrash}`);
  const dataset = tf.data.zip({ // setup data
    xs: tf.data.array([[1],[1]]),
    ys: tf.data.array([1]),
  }).batch(1);

  const model = tf.sequential({ // setup model
    layers: [
      tf.layers.dense({units: 1, inputShape: [1]}),
      tf.layers.reshape({targetShape: [(shouldCrash ? 2 : 1)]}), // use invalid shape when crashing
    ],
  });
  model.compile({ optimizer: 'sgd', loss: 'meanSquaredError' });
  console.log('  Tensors before:', tf.memory().numTensors);
  try {
    const history = await model.fitDataset(dataset, { epochs: 1 });
  } catch (err) {
    console.log(`    Error: ${err.message}`);
  }
  
  console.log('  Tensors after:', tf.memory().numTensors);
  return model
}

(async () => {
  const m1 = await train(false); // normal training
   tf.dispose(m1)
  const m2 = await train(true); // training with error
  
  tf.dispose(m2)
  tf.disposeVariables() 
  console.log('Tensors afters:', tf.memory().numTensors);
   
})();
<script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@1.1.2/dist/tf.min.js"></script>
...