时间戳

首页 > 解决方案 > 在 Keras 的 GPU 上微调 VGG-16：内存消耗

python - 在 Keras 的 GPU 上微调 VGG-16：内存消耗

问题描述

我正在为我的任务微调 VGG-16。这个想法是我加载预训练的权重，删除最后一层（这是具有 1000 个输出的 softmax）并用具有几个输出的 softmax 替换它。然后我冻结除最后一层以外的所有层并训练模型。

这是构建原始模型并加载权重的代码。

def VGG_16(weights_path=None):
    model = Sequential()
    model.add(ZeroPadding2D((1,1),input_shape=(224,224,3)))
    model.add(Conv2D(64, (3, 3), activation='relu'))
    model.add(ZeroPadding2D((1,1)))
    model.add(Conv2D(64, (3, 3), activation='relu'))
    model.add(MaxPooling2D((2,2), strides=(2,2)))

    model.add(ZeroPadding2D((1,1)))
    model.add(Conv2D(128, (3, 3), activation='relu'))
    model.add(ZeroPadding2D((1,1)))
    model.add(Conv2D(128, (3, 3), activation='relu'))
    model.add(MaxPooling2D((2,2), strides=(2,2)))

    model.add(ZeroPadding2D((1,1)))
    model.add(Conv2D(256, (3, 3), activation='relu'))
    model.add(ZeroPadding2D((1,1)))
    model.add(Conv2D(256, (3, 3), activation='relu'))
    model.add(ZeroPadding2D((1,1)))
    model.add(Conv2D(256, (3, 3), activation='relu'))
    model.add(MaxPooling2D((2,2), strides=(2,2)))

    model.add(ZeroPadding2D((1,1)))
    model.add(Conv2D(512, (3, 3), activation='relu'))
    model.add(ZeroPadding2D((1,1)))
    model.add(Conv2D(512, (3, 3), activation='relu'))
    model.add(ZeroPadding2D((1,1)))
    model.add(Conv2D(512, (3, 3), activation='relu'))
    model.add(MaxPooling2D((2,2), strides=(2,2)))

    model.add(ZeroPadding2D((1,1)))
    model.add(Conv2D(512, (3, 3), activation='relu'))
    model.add(ZeroPadding2D((1,1)))
    model.add(Conv2D(512, (3, 3), activation='relu'))
    model.add(ZeroPadding2D((1,1)))
    model.add(Conv2D(512, (3, 3), activation='relu'))
    model.add(MaxPooling2D((2,2), strides=(2,2)))

    model.add(Flatten())
    model.add(Dense(4096, activation='relu'))
    model.add(Dropout(0.5))
    model.add(Dense(4096, activation='relu'))
    model.add(Dropout(0.5))
    model.add(Dense(1000, activation='softmax'))

    if weights_path:
        model.load_weights(weights_path)

    return model

在我的案例中，Keras 使用 Tensorflow 作为后端。Tensorflow 是为使用 GPU（使用 CUDA）而构建的。我目前有一张相当旧的卡：具有 2Gb 内存的 GTX 760。

在我的卡上，由于内存不足错误，我什至无法加载整个模型（上面的代码）。

在这里，作者说 4Gb 也是不够的。

在这里，GTX 1070 甚至能够训练 VGG-16（不仅仅是将其加载到内存中），但只能使用一些批量大小和不同的框架（不在 Keras 中）。GTX 1070 似乎总是有 8Gb 的内存。

所以看起来 4Gb 显然不足以微调 VGG-16，8Gb可能就足够了。

问题是：多少内存足以用 Keras+TF 微调 VGG-16？6Gb 就足够了，或者 8Gb 是最低限度并且可以，还是需要更大的？

标签： pythontensorflowmachine-learningkerasvgg-net

解决方案

我在 Tensorflow 中微调了 VGG-16，批量大小为 32（GPU：8GB）。我认为这与 Keras 使用 Tensorflow 的情况相同。但是，如果您想使用更大的批量进行训练，则可能需要 12 或 16 GB GPU。

推荐阅读