TensorFlow内存管理问题咨询及自建网络代码说明

Hey there! Let's dive into practical TensorFlow memory management fixes and optimizations that fit right with your convolutional network code. I'll break this down into actionable steps you can implement immediately:

1. Graph & Session Configuration Tweaks #

• Dynamic GPU Memory Allocation
  By default, TensorFlow grabs all available GPU memory upfront, which can lead to unnecessary waste. Fix this by configuring your session to allocate memory on demand:

  config = tf.ConfigProto()
  config.gpu_options.allow_growth = True  # Only uses what it needs
  # Optional: Cap memory usage to a percentage (e.g., 70% of GPU)
  # config.gpu_options.per_process_gpu_memory_fraction = 0.7
  with tf.Session(config=config) as sess:
      # Run your training/inference here
  

  This is one of the quickest wins for preventing out-of-memory (OOM) errors.

• Clean Up Old Graphs
  If you're calling _create_network multiple times (e.g., for target vs online networks in DQN), always reset the default graph first to avoid redundant nodes piling up:

  tf.reset_default_graph()
  self._create_network(scope='network')
  

2. Variable & Layer Optimization #

• Reuse Variables Across Scopes
  Your code uses variable_scope, which is perfect for reusing weights (critical for architectures like DQN). Add a reuse parameter to your network function to avoid duplicate variables that bloat memory:

  def _create_network(self, scope='network', reuse=False):
      with tf.variable_scope(scope, reuse=reuse):
          self.inputs = tf.placeholder(shape=[None, *self.img_shape, self.hist_len], dtype=tf.float32)
          self.conv_1 = slim.conv2d(activation_fn=tf.nn.relu, inputs=self.inputs, num_outputs=16, kernel_size=[8, 8], stride=4, padding='SAME')
          self.conv_2 = slim.conv2d(activation_fn=tf.nn.relu, inputs=self.conv_1, num_outputs=64, kernel_size=[4, 4], stride=2, padding='SAME')
          self.fc = slim.fully_connected(...)
  

  Then reuse it like this:

  # Build online network
  self._create_network(scope='online_net')
  # Build target network with reused weights
  self._create_network(scope='target_net', reuse=True)
  

• Avoid Unnecessary Tensor Retention
  You don't need to assign every layer to a self attribute unless you explicitly need to reference it later. Use local variables for intermediate layers to let TensorFlow garbage-collect unused tensors:

  with tf.variable_scope(scope, reuse=reuse):
      inputs = tf.placeholder(shape=[None, *self.img_shape, self.hist_len], dtype=tf.float32)
      conv_1 = slim.conv2d(activation_fn=tf.nn.relu, inputs=inputs, num_outputs=16, kernel_size=[8, 8], stride=4, padding='SAME')
      conv_2 = slim.conv2d(activation_fn=tf.nn.relu, inputs=conv_1, num_outputs=64, kernel_size=[4, 4], stride=2, padding='SAME')
      self.fc = slim.fully_connected(conv_2, ...)  # Only keep what you need
  

3. Batch & Shape Adjustments #

• Tune Batch Size
  Your placeholder uses a dynamic batch size (None), which is flexible, but a batch that's too large will eat up GPU memory. Start with a smaller batch size (e.g., 32 or 16) and scale up only if your GPU has leftover memory.

• Use Static Shapes Where Safe
  If your input shape is fixed during training, replace *self.img_shape with explicit values (e.g., 84,84) instead of relying on dynamic shapes. TensorFlow can optimize memory allocation better with static shape information.

4. Debugging Memory Leaks #

• Profile Memory Usage
  Use TensorFlow's built-in profiler to identify which layers or variables are hogging memory:

  from tensorflow.python.profiler import model_analyzer
  
  # After building your graph
  profiler = model_analyzer.Profiler(graph=tf.get_default_graph())
  profiler.profile_name_scope(
      options=model_analyzer.ProfileOptionBuilder().with_memory().build()
  )
  

  This will show you a breakdown of memory usage per layer, making it easy to spot bottlenecks.

• Monitor Real-Time GPU Usage
  Run tf.contrib.memory_stats.MaxBytesInUse() during session runs to track peak memory usage:

  max_mem = tf.contrib.memory_stats.MaxBytesInUse()
  print("Peak GPU memory used:", sess.run(max_mem))
  

内容的提问来源于stack exchange，提问作者GoingMyWay