如何减少GPU上的内存使用量Tensorflow

时间:2019-07-03 00:55:42

标签: python tensorflow memory gpu

我正在使用Tensorflow进行科学计算,将其作为GPU自动区分的框架,而不是机器学习的框架。我的代码可在CPU和GPU上运行,但是一旦问题大小增加,GPU版本就会失败。我怀疑GPU的内存不足,因此我收到此错误。我的代码在GPU上的运行速度至少快10倍,所以我宁愿在GPU上运行它。我想要一些建议来修改我的代码或设置,以便我可以继续使用GPU运行Tensorflow代码。

我本机上同时编译了tf 2.0.0-beta1的cpu和gpu版本。 2个Intel®Xeon®CPU E5-2680 v4,具有512 GB内存, 泰坦V 12GB。 我已经使用了tensorboard profiler和图形显示来尝试找出内存使用情况,但是没有找到有用的信息。 我发现的只是我的代码在Adam优化的第一个周期的中间崩溃。 For context this is what my graph looks like

@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]]

我也尝试在tf.config.optimizer.set_experimental_options中禁用ArithmeticOptimizer选项,但是却出现了matrix_exponential / ArithmeticOptimizer / ADDN错误。

0 个答案:

没有答案