如何计算TensorFlow中训练单样本CNN所需的GPU内存

我正尝试计算特定数据集上训练CNN所需的GPU内存（或通用计算方法），因为训练时出现显存不足，但无法定位原因。

实现的CNN模型 #

我用TensorFlow实现了如下CNN：

from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv2D, Dense, Flatten

model = Sequential()
model.add(Conv2D(16, (3,3), activation='relu', input_shape=(1491, 257,1)))
model.add(Conv2D(16, (3,3), activation='relu'))
model.add(Flatten())
model.add(Dense(128, activation='relu'))
model.add(Dense(1, activation='sigmoid'))

model.summary()

模型输出的summary信息：

Model: "sequential_1"
_________________________________________________________________
 Layer (type)                Output Shape              Param #   
=================================================================
 conv2d_2 (Conv2D)           (None, 1489, 255, 16)     160       
                                                                 
 conv2d_3 (Conv2D)           (None, 1487, 253, 16)     2320      
                                                                 
 flatten_1 (Flatten)         (None, 6019376)           0         
                                                                 
 dense_1 (Dense)             (None, 128)               770480256 
                                                                 
 dense_2 (Dense)             (None, 1)                 129       
                                                                 
=================================================================
Total params: 770482865 (2.87 GB)
Trainable params: 770482865 (2.87 GB)
Non-trainable params: 0 (0.00 Byte)

存储该模型需要2.87GB内存。

单样本训练内存计算 #

网络输入为1491x257x1张量，输入来自以下数据管道处理的wav文件：

data = data.map(preprocess)
data = data.cache()
data = data.shuffle(buffer_size=1000)
data = data.batch(1)
data = data.prefetch(1)

train = data.take(1)
test = data.skip(10).take(1)

preprocess函数会将每个wav文件转换为1491x257x1的频谱图（函数详情见文末）。

根据CNN内存消耗计算逻辑，单个数据元素前向传播所需内存为输入元素内存加上各层输出内存：

• 单个输入元素内存：4*(1491*257*1) = 1,532,748字节（TensorFlow默认float32类型，占4字节）
• 各层输出内存（Flatten仅重排数据，不计额外内存）：4 * 输出维度

由此计算单样本前向传播所需内存：

4*(1491*257*1 + 1489*255*16 + 1487*253*16 + 128 + 1) Bytes

换算为GB：

4*(1491*257*1 + 1489*255*16 + 1487*253*16 + 128 + 1)/(1024**3) GB = 0.05 GB

单样本训练总GPU内存预估：模型内存+前向传播内存+反向传播内存（假设与前向一致），即2.87GB + 0.05GB + 0.05GB = 2.97GB。

数据集内存估算 #

我复现的是鸟类音频分类项目，原始文件占用0.4 GB内存。预处理后每个文件转为float32类型的1491x257x1频谱图：

• 810个频谱图总内存：810 * 4 * (1491 * 257)/(1024 * 1024 * 1024) = 1.15GB
• 数据管道设置了1000个元素的shuffle缓冲区，预估占用：1000 * 4 * (1491 * 257)/(1024 * 1024 * 1024) = 1.43GB

当前核算的最大内存占用为2.97GB + 0.40 GB + 1.43GB = 4.80GB，而我的GPU有9.29GB内存，剩余4.49GB，理论上单样本训练不应出现显存不足。我忽略了什么？

预处理函数详情 #

import tensorflow as tf 
import tensorflow_io as tfio

def preprocess(file_path, label): 
    wav = load_wav_16k_mono(file_path)
    wav = wav[:48000]
    zero_padding = tf.zeros([48000] - tf.shape(wav), dtype=tf.float32)
    wav = tf.concat([zero_padding, wav],0)
    spectrogram = tf.signal.stft(wav, frame_length=320, frame_step=32)
    spectrogram = tf.abs(spectrogram)
    spectrogram = tf.expand_dims(spectrogram, axis=2)
    return spectrogram, label

def load_wav_16k_mono(filename):
    # Load encoded wav file
    file_contents = tf.io.read_file(filename)
    # Decode wav (tensors by channels) 
    wav, sample_rate = tf.audio.decode_wav(file_contents, desired_channels=1)
    # Removes trailing axis
    wav = tf.squeeze(wav, axis=-1)
    sample_rate = tf.cast(sample_rate, dtype=tf.int64)
    # Goes from 44100Hz to 16000hz - amplitude of the audio signal
    wav = tfio.audio.resample(wav, rate_in=sample_rate, rate_out=16000)
    return wav

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