python - 如何找出模型本身使用的 VRAM 数量？(LSTM)

Question

如何找出此模型的 VRAM 使用情况？（这不是关于正在训练的数据，而是关于加载到 VRAM 中的模型及其权重

from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense,Dropout,LSTM,BatchNormalization
import tensorflow as tf


model = Sequential()
model.add(LSTM(700, input_shape=(10000,5000,20), return_sequences=True))
model.add(Dropout(0.2))
model.add(BatchNormalization())


model.add(LSTM(700, return_sequences=True))
model.add(Dropout(0.2))
model.add(BatchNormalization())


model.add(LSTM(700))
model.add(Dropout(0.2))
model.add(BatchNormalization())


model.add(Dense(64, activation='sigmoid'))
model.add(Dropout(0.2))


model.add(Dense(32, activation='sigmoid'))
model.add(Dropout(0.2))

Answer 1

您可以根据参数的数量和参数数据类型对其进行近似。

model.summary()

会告诉你有多少参数，然后乘以大小。例如，如果你发现你有 5000 个参数并且它们是 32b 浮点，那么你的模型至少是 5000*32/8 = 1250 字节。

当然，还有更多内容，但这通常是获得相当接近的下限的简单方法。您还可以保存模型并查看保存的文件有多大，其中还包括模型结构。

tf.saved_model.save(model, path)

一定要小心，因为如果你的模型很大，那么训练它会更大，因为它也会存储所有可学习参数的梯度。所以你可以很容易地想象一个 1GB 的模型可能需要 5-10GB 来训练。