具有Tensorflow(1.3)后端的Keras(2.0.8)获取所有可用RAM

时间:2018-03-07 12:44:51

标签: python tensorflow keras ram

我正在使用带有keras后端的tensorflow库并启用了CUDA。请参阅PIP包版本输出:

Keras (2.0.8)
tensorflow-gpu (1.3.0)
tensorflow-tensorboard (0.1.8)

我有以下代码创建VGG16模型并加载ImageNet权重:

def create_vgg16_model(target_size: tuple, n_classes: int):
    base = VGG16(include_top=False,
                 input_shape=target_size,
                 weights='imagenet')

    x = base.output
    x = Flatten()(x)
    x = Dense(n_classes, activation='softmax', name='top')(x)

    model = Model(inputs=base.input, outputs=x)
    for layer in model.layers[:-1]:
        layer.trainable = False

    model.compile(optimizer='adam', loss='categorical_crossentropy')
    return model

该模型的培训进展顺利,nvidia-smi表明根据需要使用了GPU内存。但后来我检查了top命令的输出,这就是我所看到的:

  PID USER      PR  NI    VIRT    RES    SHR S  %CPU %MEM     TIME+ COMMAND                                                                 
 1268 ck        20   0  166288  31964  12416 S  29.5  0.1  13:05.39 Xtightvnc                                                               
32235 ck        30  10   32252   3700   3356 S   5.3  0.0   0:36.93 cwaves 
------------------------------------------------------------------------------    
32212 ck        20   0 27.485g 1.184g 190404 S   2.3  3.8   0:35.44 python  
------------------------------------------------------------------------------                                                                
26015 root      20   0       0      0      0 S   0.3  0.0   0:00.30 kworker/3:1                                                             
31754 ck        20   0   43168   3904   3080 R   0.3  0.0   0:04.45 top                                                                     
    1 root      20   0  185644   6204   3984 S   0.0  0.0   0:10.44 systemd                                                                 

我已经使用调试器完成了代码,并意识到内存是在keras.backend.tensorflow_backed中创建tf.Session对象的以下函数中分配的:

def get_session():        
    global _SESSION
    if tf.get_default_session() is not None:
        session = tf.get_default_session()
    else:
        if _SESSION is None:
            if not os.environ.get('OMP_NUM_THREADS'):
                config = tf.ConfigProto(allow_soft_placement=True)
            else:
                num_thread = int(os.environ.get('OMP_NUM_THREADS'))
                config = tf.ConfigProto(intra_op_parallelism_threads=num_thread,
                                        allow_soft_placement=True)
                # next line allocates ~28GB of RAM
                _SESSION = tf.Session(config=config)
        session = _SESSION
    if not _MANUAL_VAR_INIT:
        with session.graph.as_default():
            _initialize_variables()
    return session

并且,所有可用模型都会发生这种情况,因为在创建会话之前,在训练开始或变量初始化之前会分配内存。

我知道TF会分配所有可用的GPU内存(除非你覆盖ConfigProto和/或调整你的环境变量),但它是否对RAM做同样的事情?即似乎框架正在分配我机器上的所有RAM,除了已经由其他进程分配的RAM。

是否有人发现此类行为有不同版本的tensorflowkeras?你认为有办法以某种方式限制已用内存的数量吗?

更新1

前段时间我的一个训练脚本被内核杀死,在50-60个训练时期后出现内存不足错误。虽然易失性GPU内存使用情况统计显示它也被使用。 (不仅仅是分配,正如我所理解的那样)。

更新2

同意,虚拟内存不是有效的指标。但我已经发现,在模型培训过程中,内存消耗几乎呈线性增长。我有以下培训循环:

def train_model(model, x, y):
    loss = model.train_on_batch(x, y)
    return loss


def train_model_42(model, x, y):
    # dummy function
    return 42.0


def training_loop():
    # training parameters
    target_size = 224, 224, 3
    batch_size = 128

    # generator yielding batches of file paths
    files_stream = FilesStream(folder=TRAIN_IMAGES, batch_size=batch_size)
    files_source = files_stream()

    # list of generators loading images from persistent storage
    gens = [
        image_loader(),
        augment_images(horizontal_flip=True),
        shuffle_samples(),
        normalize_images(target_size=target_size)
    ]

    # Model from keras.applications with replaced top layer
    model = get_model('resnet50').build(n_classes=n_classes)

    for epoch in range(1, 1001):
        epoch_loss = []
        for _ in range(files_stream.steps_per_epoch):
            for gen in gens:
                gen.send(None)
            processed = next(files_source)
            for gen in gens:
                processed = gen.send(processed)
            x, y = processed
            loss = train_model_42(model, x, y) # <-- this shows pic. 1
            # loss = train_model(model, x, y)    <-- this shows pic. 2                  
            epoch_loss.append(loss)
        avg_loss = sum(epoch_loss) / len(epoch_loss)
        print('Epoch %03d: train loss = %2.4f' % (epoch, avg_loss))

当我使用虚拟训练功能时,内存消耗情节看起来好像显示在pic 1上: enter image description here

但是在运行真正的培训过程时,它看起来像pic 2enter image description here

为什么在培训过程中内存消耗会增加?以前批量的数据是否已缓存?模型/重量或其他什么东西会占用越来越多的记忆吗?

我认为我的数据预处理管道可能有问题,但我故意将预处理功能写成生成器。是否可以将某种默认的Keras回调应用于跟踪负责增加内存使用量的培训信息的模型?

1 个答案:

答案 0 :(得分:1)

我想我已经找到了问题的根源。果然,它与tensorflowkeras无关,而是我使用它们的方法。

这是一个类似于我的图像增强功能的功能:

def augment_images():
    transformer = ImageDataGenerator()
    while True:
        x, y = yield
        generator = transformer.flow(x, y, batch_size=len(x), shuffle=False)
        transformed = next(generator)
        yield transformed

它使用keras.preprocessing.image.ImageDataGenerator类来增强图像。但是该类本身实例化NumpyArrayIterator对象,其中将引用保持为xy批处理并将ImageDataGenerator作为委托进行调用。而且,这是内存泄漏的根源。似乎这些对象阻止了数组被垃圾收集。

这是一个更新的扩充函数,它明确地使用了迭代器:

def augment_images(width_shift=0.2,
                   height_shift=0.2,
                   zoom=0.2,
                   rotation=30,
                   vertical_flip=False,
                   horizontal_flip=False):

    transformer = ImageDataGenerator()
    iterator = None

    while True:
        x, y = yield
        if iterator is None:
            iterator = NumpyArrayIterator(
                x, y, transformer,
                batch_size=len(x),
                shuffle=False,
                seed=None,
                data_format=transformer.data_format)
        else:
            iterator.n = x.shape[0]
            iterator.x = x
            iterator.y = y
        transformed = next(iterator)
        yield transformed

所以,问题出在我用来预处理数据的生成器包装器中。 (或者我会说,在我使用Keras的API和Python的生成器的方法中)。至少现在,当我取代图像增强功能时,不再有内存泄漏。