如何计算使用 Tensorflow 在单个数据元素上训练 CNN 所需的内存

Question

我正在尝试计算在特定数据集（或者一般情况下）上训练特定 CNN 所需的 GPU 内存量。原因是当我尝试训练这个网络时，我的内存不足，但无法弄清楚为什么。

我在 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()

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

频谱图图像（为了彻底起见，我将在本问题的底部包含

preprocess

及其调用的函数）。

按照计算 CNN 内存消耗中的解释，转发单个数据元素所需的内存量是存储该元素所需的内存量加上网络每层输出所需的内存量。

我相信单个数据元素占用的内存是

4*(1491*257*1) = 1,532,748 bytes

，因为我的

input_shape=(1491,257,1)

和我目前使用的是Tensorflow的默认类型

float32

，它有4个字节。同样，存储每层输出所需的数据量（除了展平，它只是前一层输出的重新排列）将是

4 * output dimensions

。这给出了向前传播单个数据元素所需的内存：

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

或

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

因此，训练单个图像将需要

2.87GB + 0.05GB = 2.92GB

的 GPU 总内存，或者如果反向支撑单个图像需要与前向支撑它相同的内存量，则需要

= 2.97GB

。

接下来，我想考虑可能用于存储 GPU 上所有数据的内存量。

我正在尝试使用 Tensorflow 重现笔记本和 Kaggle 鸟类音频分类的结果（Kaggle Bird Audio Classification）。使用 df -h，我可以测量原始文件所需的内存为

0.4 GB

。同样，由于预处理步骤将每个文件转换为 float32 的

1491x257x1

频谱图，并且有 810 个频谱图，我知道存储整组频谱图需要

810 * 4 * (1491 * 257)/(1024 * 1024 * 1024) = 1.15GB

或者因为我有 1000 个数据元素的缓冲区，tensorflow 可能会觉得它需要

1000 * 4 * (1491 * 257)/(1024 * 1024 * 1024) = 1.43GB

我现在最多可以解释

2.97GB + 0.40 GB + 1.43GB = 4.80GB

自

9.29GB - 4.80GB = 4.49GB

以来，我认为我仍然应该有足够的内存来训练单个图像，而不会遇到内存不足错误。我在这里缺少什么？

---

创建 TF 数据集所使用的函数详细信息：

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

Answer 1

您可以使用

tf.data.AUTOTUNE

功能进行自动记忆拟合。 （参考-> https://www.tensorflow.org/guide/data_performance#prefetching）

也许您会对我的鸟声分类笔记本感兴趣 (76 GB) 160 种波兰鸟声分类，带集成深度学习

你应该尝试；

data = data.map(preprocess, num_parallel_calls = tf.data.AUTOTUNE)
data = data.cache()
data = data.shuffle(buffer_size=1000)
data = data.batch(1)
data = data.prefetch(tf.data.AUTOTUNE)