TensorFlow GradientTape.jacobian() GPU显存不足，双RTX3090如何利用？

核心问题分析 #

计算56×56矩阵对模型可训练变量的Jacobian时，单GPU无法承载超大尺寸的梯度张量（报错中的[3136,56,8,8,64]张量），可通过双GPU显存拆分和Jacobian计算优化两类方案解决。

一、利用双GPU扩容显存的具体方案 #

1. 数据并行拆分Jacobian计算任务 #

将56×56矩阵的元素拆分到两个GPU上分别计算梯度，最后合并结果。以下是手动指定设备分片的实现示例：

import tensorflow.compat.v2 as tf
import gc

# 保留原模型定义
class NN(tf.keras.Model):
  """The convolutional network used to compute the agent's policy."""
  def __init__(self,
               num_actions,
               name = None,
               dropout = 0.0,
               rand_conv = False,
               projection = True,
               **kwargs):
    super(NN, self).__init__(**kwargs)
    self.kwargs = kwargs
    self._dropout = dropout
    self._num_actions = num_actions
    self.rand_conv = None
    self._projection = projection
    activation_fn = tf.keras.activations.relu
    self.conv0 = tf.keras.layers.Conv2D(
        32, [8, 8],
        strides=4,
        padding='same',
        activation=activation_fn,
        name='Conv')
    self.conv1 = tf.keras.layers.Conv2D(
        64, [4, 4],
        strides=2,
        padding='same',
        activation=activation_fn,
        name='Conv1')
    self.conv2 = tf.keras.layers.Conv2D(
        64, [3, 3],
        strides=1,
        padding='same',
        activation=activation_fn,
        name='Conv2')
    self.flatten = tf.keras.layers.Flatten()
    self.dense0 = tf.keras.layers.Dense(256, activation=activation_fn)
    self.dense1 = tf.keras.layers.Dense(64, activation=activation_fn)
    self.dense2 = tf.keras.layers.Dense(num_actions, name='fully_connected')

  @tf.function
  def call(self, state, training=True):
    x = self.representation(state, projection=False)
    if training:
      x = tf.nn.dropout(x, rate=self._dropout)
    x = self.dense2(x)
    if self._num_actions == 1:
      x = tf.squeeze(x, axis=-1)
    return x

  def representation(self, state, projection=True):
    x = tf.cast(state, tf.float32)
    if self.rand_conv is not None:
      x = self.rand_conv(x)
    x = self.conv0(x)
    x = self.conv1(x)
    x = self.conv2(x)
    x = self.flatten(x)
    x = self.dense0(x)
    if projection and self._projection:
      x = self.dense1(x)
    return x

nn_model = NN(num_actions=2, dropout=0.0, rand_conv=True, projection=True)
fake1 = tf.random.uniform([56,60,60,2])
fake2 = tf.random.uniform([56,60,60,2])

# 拆分矩阵为前后两部分，分配到不同GPU计算
split_idx = 28
jacobian_parts = []

# GPU0计算前28行的Jacobian
with tf.device('/GPU:0'):
    with tf.GradientTape(persistent=True) as tape0:
        tape0.watch(nn_model.trainable_variables)
        out1 = nn_model.representation(fake1)
        out2 = nn_model.representation(fake2)
        cos = tf.keras.losses.CosineSimilarity(axis=2, reduction=tf.keras.losses.Reduction.NONE)
        m_part0 = cos(tf.expand_dims(out1[:split_idx], 1), tf.expand_dims(out2, 0))
    jac_part0 = tape0.jacobian(m_part0, nn_model.trainable_variables)
    jacobian_parts.append(jac_part0)
    del tape0, m_part0
    gc.collect()

# GPU1计算后28行的Jacobian
with tf.device('/GPU:1'):
    with tf.GradientTape(persistent=True) as tape1:
        tape1.watch(nn_model.trainable_variables)
        out1 = nn_model.representation(fake1)
        out2 = nn_model.representation(fake2)
        cos = tf.keras.losses.CosineSimilarity(axis=2, reduction=tf.keras.losses.Reduction.NONE)
        m_part1 = cos(tf.expand_dims(out1[split_idx:], 1), tf.expand_dims(out2, 0))
    jac_part1 = tape1.jacobian(m_part1, nn_model.trainable_variables)
    jacobian_parts.append(jac_part1)
    del tape1, m_part1
    gc.collect()

# 合并两部分结果
full_jacobian = []
for j0, j1 in zip(jacobian_parts[0], jacobian_parts[1]):
    full_jac = tf.concat([j0, j1], axis=0)
    full_jacobian.append(full_jac)

2. 模型并行：将模型层分配到不同GPU #

把模型的不同层拆分到两个GPU上，让模型参数和中间计算结果分散到两张卡的显存中，降低单卡负载：

class NN(tf.keras.Model):
    def __init__(self, num_actions, **kwargs):
        super(NN, self).__init__(**kwargs)
        self._dropout = 0.0
        self._num_actions = num_actions
        self.rand_conv = None
        self._projection = True
        activation_fn = tf.keras.activations.relu
        
        # 卷积层分配到GPU0
        with tf.device('/GPU:0'):
            self.conv0 = tf.keras.layers.Conv2D(32, [8,8], strides=4, padding='same', activation=activation_fn)
            self.conv1 = tf.keras.layers.Conv2D(64, [4,4], strides=2, padding='same', activation=activation_fn)
        # 后续层分配到GPU1
        with tf.device('/GPU:1'):
            self.conv2 = tf.keras.layers.Conv2D(64, [3,3], strides=1, padding='same', activation=activation_fn)
            self.flatten = tf.keras.layers.Flatten()
            self.dense0 = tf.keras.layers.Dense(256, activation=activation_fn)
            self.dense1 = tf.keras.layers.Dense(64, activation=activation_fn)
            self.dense2 = tf.keras.layers.Dense(num_actions)
    
    # 原call和representation方法保持不变

二、Jacobian计算的显存优化（配合双GPU效果更佳） #

1. 逐元素/逐批次计算Jacobian #

避免一次性计算整个矩阵的Jacobian，循环处理小批次元素，每次只占用部分显存：

jacobian_list = []
# 逐行计算Jacobian
for i in range(56):
    with tf.GradientTape() as tape:
        tape.watch(nn_model.trainable_variables)
        out1 = nn_model.representation(fake1)
        out2 = nn_model.representation(fake2)
        cos = tf.keras.losses.CosineSimilarity(axis=2, reduction=tf.keras.losses.Reduction.NONE)
        m = cos(tf.expand_dims(out1, 1), tf.expand_dims(out2, 0))
        row_element = m[i]
    jac = tape.jacobian(row_element, nn_model.trainable_variables)
    jacobian_list.append(jac)
# 整理为完整Jacobian张量
full_jacobian = tf.stack(jacobian_list, axis=0)

2. 启用混合精度训练 #

将部分张量转为float16格式，大幅降低显存占用，同时不影响计算精度：

tf.keras.mixed_precision.set_global_policy('mixed_float16')
# 之后初始化模型并执行计算流程

三、额外显存清理技巧 #

• 每次计算后手动删除无用张量，并调用gc.collect()释放内存
• 禁用TensorFlow显存预分配，让显存按需使用：

  gpus = tf.config.list_physical_devices('GPU')
  for gpu in gpus:
      tf.config.experimental.set_memory_growth(gpu, True)
  

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