Python 如何确定模型本身使用的VRAM数量？（LSTM）

如何了解该型号的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))

---

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

model.summary()

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

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

tf.saved_model.save(model, path)

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