Tensorflow：6 层 CNN：OOM（使用 10Gb GPU 内存）

Question

我正在使用以下代码运行 6 层 CNN，顶部有 2 个 FC 层（在 Tesla K-80 GPU 上）。 不知何故，它消耗了整个内存 10GB 并且内存不足。我知道我可以减少 batch_size 然后运行，但我也想运行 15 或 20 个 CNN 层。下面的代码有什么问题以及为什么它需要所有记忆？我应该如何运行 15 层 CNN 的代码。

代码：

import model


with tf.Graph().as_default() as g_train:
    filenames = tf.train.match_filenames_once(FLAGS.train_dir+'*.tfrecords')
    filename_queue = tf.train.string_input_producer(filenames, shuffle=True, num_epochs=FLAGS.num_epochs)
    feats,labels = get_batch_input(filename_queue, batch_size=FLAGS.batch_size)
    ### feats size=(batch_size, 100, 50)
    logits = model.inference(feats, FLAGS.batch_size)
    loss = model.loss(logits, labels, feats)
    tvars = tf.trainable_variables()
    global_step = tf.Variable(0, name='global_step', trainable=False)
    # Add to the Graph operations that train the model.
    train_op = model.training(loss, tvars, global_step, FLAGS.learning_rate, FLAGS.clip_gradients)
    # Add the Op to compare the logits to the labels during evaluation.
    eval_correct = model.evaluation(logits, labels, feats)
    summary_op = tf.merge_all_summaries()
    saver = tf.train.Saver(tf.all_variables(), max_to_keep=15)

    # The op for initializing the variables.
    init_op = tf.initialize_all_variables()

    sess = tf.Session()
    sess.run(init_op)

    summary_writer = tf.train.SummaryWriter(FLAGS.model_dir,
                                        graph=sess.graph)

    # Start input enqueue threads.
    coord = tf.train.Coordinator()
    threads = tf.train.start_queue_runners(sess=sess, coord=coord)
    try:
        step = 0
        while not coord.should_stop():
            _, loss_value = sess.run([train_op, loss])

            if step % 100 == 0:
               print('Step %d: loss = %.2f (%.3f sec)' % (step, loss_value))
               # Update the events file.
               summary_str = sess.run(summary_op)
               summary_writer.add_summary(summary_str, step)

            if (step == 0) or (step + 1) % 1000 == 0 or (step + 1) == FLAGS.max_steps:
               ckpt_model = os.path.join(FLAGS.model_dir, 'model.ckpt')
               saver.save(sess, ckpt_model, global_step=step)
               #saver.save(sess, FLAGS.model_dir, global_step=step)

        step += 1
    except tf.errors.OutOfRangeError:
        print('Done training for %d epochs, %d steps.' % (FLAGS.num_epochs, step))
    finally:
        coord.join(threads)
        sess.close()

###################### File model.py #################### 
def conv2d(x, W, b, strides=1):
    # Conv2D wrapper, with bias and relu activation
    x = tf.nn.conv2d(x, W, strides=[1, strides, strides, 1],     
            padding='SAME')
    x = tf.nn.bias_add(x, b)
    return tf.nn.relu(x)

def maxpool2d(x, k=2,s=2):
    # MaxPool2D wrapper
    return tf.nn.max_pool(x, ksize=[1, k, k, 1], strides=[1, s,  
    s,1],padding='SAME')


def inference(feats,batch_size):
    #feats size (batch_size,100,50,1) #batch_size=256
    conv1_w=tf.get_variable("conv1_w",            [filter_size,filter_size,1,256],initializer=tf.uniform_unit_scaling_initializer())
    conv1_b=tf.get_variable("conv1_b",[256])
    conv1 = conv2d(feats, conv1_w, conv1_b,2)
    conv1 = maxpool2d(conv1, k=2,s=2)

    ### This was replicated for 6 layers and the 2 FC connected layers are added

    return logits

def training(loss, train_vars, global_step, learning_rate, clip_gradients):
  # Add a scalar summary for the snapshot loss.
  tf.scalar_summary(loss.op.name, loss)
  grads, _ = tf.clip_by_global_norm(tf.gradients(loss,  train_vars,aggregation_method=1), clip_gradients)
  optimizer = tf.train.AdamOptimizer(learning_rate)
  train_op = optimizer.apply_gradients(zip(grads, train_vars),  global_step=global_step)
  return train_op

Answer 1

我不太确定模型 python 库是什么。如果它是您编写的并且可以更改优化器中的设置，我会建议在我自己的代码中使用以下内容

train_step = tf.train.AdamOptimizer(learning_rate).minimize(cost, aggregation_method = tf.AggregationMethod.EXPERIMENTAL_ACCUMULATE_N)

默认情况下 aggeragetion_method 是 ADD_N 但如果你将其更改为 EXPERIMENTAL_ACCUMULATE_N 或 EXPERIMENTAL_TREE 这将大大节省内存。这些程序的主要内存消耗是 tensorflow 必须保存每个神经元的输出值，以便它可以计算梯度。根据我的经验，更改aggregation_method 有很大帮助。

顺便说一句，我认为您的代码没有任何问题。我也会在小型 cov-net 上耗尽内存。