训练 Keras 模型需要多少 VRAM

Question

您好，我最近开始了我的学士项目，其中我需要训练一个具有以下结构的 LSTM 模型

输入（输入层）[（无，无，80）] 0

掩蔽（掩蔽）（无、无、80）0

lstm（LSTM）（无，无，100）72400

输出（TimeDistributed）（无，无，39）3939

=================================================== ================

总参数：76,339

可训练参数：76,339

不可训练的参数：0

inputs = tf.keras.Input(shape=(None,nb_features), name = 'inputs')
    
    x = tf.keras.layers.Masking(mask_value = data.MASK_VALUE)(inputs)
    
    x = tf.keras.layers.LSTM(hidden_units,
                             return_sequences = True,
                             dropout = dropout_rate)(x)

    dense = tf.keras.layers.Dense(nb_skills, activation = 'sigmoid')
    outputs = tf.keras.layers.TimeDistributed(dense, name = 'outputs')(x)

现在我想知道我需要多少 VRAM 来进行训练。我的工作站内有一台带有 11GB VRAM 的 2080ti，但我不知道这是否足够。也许有一些有用的网站可以计算类似的东西。我尝试自己找到它，但我认为这样的东西还不存在。

感谢您的帮助

Answer 1

11GB VRAM 应该足够了。我训练了一个具有 3 个 LSTM 层的 400K+ 参数模型，可以轻松放入我的 4GB VRAM 中。如果使用 GPU，您还可以在开始训练模型时使用 Tensorboard 分析来查看模型的内存配置文件。您可以在其中看到 GPU 设备上的峰值 VRAM 使用情况。

Answer 2

您可以通过以下方式计算任何模型所需的 VRAM：https://howmuchvram.com/