Как уменьшить использование памяти Tensorflow на GPU - PullRequest
Новая
       18

Как уменьшить использование памяти Tensorflow на GPU

0 голосов
/ 03 июля 2019

Я использую Tensorflow для научных вычислений в качестве основы для автоматической дифференциации с помощью графического процессора, а не для машинного обучения. Мой код работает на CPU и GPU, но, как только размер проблемы увеличивается, версия GPU дает сбой. Я подозреваю, что GPU не хватает памяти, поэтому я получаю эту ошибку. Мой код работает как минимум в 10 раз быстрее на GPU, поэтому я бы предпочел запустить его на GPU. Я хотел бы получить совет по изменению моего кода или настроек, чтобы я мог продолжить работу с кодом Tensorflow на графическом процессоре.

Я скомпилировал версии tf 2.0.0-beta1 для cpu и gpu на своей машине. 2 процессора Intel® Xeon®® E5-2680 v4 с оперативной памятью 512 ГБ, Титан V 12ГБ. Я использовал профилировщик тензорной доски и графическое отображение, чтобы попытаться выяснить использование памяти, но не нашел полезной информации. Все, что я нашел, - это мой код, который вылетает в середине первого цикла оптимизации Адама. Для контекста это то, как выглядит мой график 

@tf.function
def optimization_step():
  with tf.GradientTape() as tape:

    tfid =mymod.proploss(mymod.trotterization(mymod.parameterization()))

    curloss = tfid 
  gradients = tape.gradient(curloss,[mymod.optimization_params])
  optimizer.apply_gradients(zip(gradients, [mymod.optimization_params]))
  return curloss



steps = range(5)
stamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
logdir = 'logs/func/%s' % stamp
writer = tf.compat.v2.summary.create_file_writer(logdir)

for step in steps:  
  tf.compat.v2.summary.trace_on(graph=True,profiler=True)
  infid = optimization_step()

  with writer.as_default():
      tf.compat.v2.summary.trace_export(
      name="my_func_trace",
      step=0,
      profiler_outdir=logdir)
  print(step,infid)

Мой код, особенно mymod.trotterization (), это в основном матричное умножение (tf.matmul) и матричные экспоненты (tf.linalg.expm). В моем коде нет явных циклов или операторов потока управления.

Я пытаюсь оптимизировать mymod.optimization_params. Некоторые из используемых тензоров: 100000x16x16 tf.complex128.

Это сообщение об ошибке, которое я получаю 

E tensorflow/core/common_runtime/executor.cc:641] Executor failed to create kernel. Not found: No registered 'TemporaryVariable' OpKernel for GPU devices compatible with node {{node StatefulPartitionedCall_1/matrix_exponential/ArithmeticOptimizer/AddOpsRewrite_add_3/tmp_var}}
     (OpKernel was found, but attributes didn't match) Requested Attributes: dtype=DT_COMPLEX128, shape=[96000,16,16], var_name="StatefulPa..._3/tmp_var", _device="/job:localhost/replica:0/task:0/device:GPU:0"
    .  Registered:  device='GPU'; dtype in [DT_INT64]
  device='GPU'; dtype in [DT_DOUBLE]
  device='GPU'; dtype in [DT_FLOAT]
  device='GPU'; dtype in [DT_HALF]
  device='CPU'

     [[StatefulPartitionedCall_1/matrix_exponential/ArithmeticOptimizer/AddOpsRewrite_add_3/tmp_var]]

Я также пытался отключить опцию ArithmeticOptimizer в tf.config.optimizer.set_experimental_options, но вместо этого я получаю ошибку matrix_exponential / ArithmeticOptimizer / ADDN.

...