J4m3s
J4m3s

Reputation: 43

Keras fit_generator using a lot of memory even with small batch sizes

Previously I manually trained my model using model.fit() inside a for loop to train it on small batches of data, due to memory constraints. The problem with this is that I can't have access to all previous histories through history.history, because it's like each time a new model is trained, and previous histories aren't stored anywhere.

When I use model.fit() on a 500 batch size, around 7 GB of my ram gets full. I use keras with tensorflow-cpu back end. But when I use a generator, even with a batch size of 50 won't fit in memory, and gets swapped onto the disk.

I'm performing classification, using 224*224 images, and I am trying to fine tune vgg face. I'm using vgg face implemented according to this link: VGG-Face

I'm using ResNet and SeNet architectures, as described in the link.

I've previously shuffled my data. I've put aside %20 of my data for test. My data, image addresses and labels, are stored in a list. The %20 of my training data will be used for validation. For example if batch size is equal to 50, train_data_generator will create a batch with size 40 from the first %80 portion of training data, and vl_data_generator will create a batch with size 10 from the last %20 portion of training data. I've written a class, and by creating an instance and invoking train method through it, I perform training. Here are generator and training parts of my code, excluding model definitions:

def prepare_input_data(self, batch_addresses):
    image = []
    for j in range(len(batch_addresses)):
        img = cv2.imread(batch_addresses[j])
        img = cv2.resize(img, (224, 224))
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

        img = img - np.array([103.939, 116.779, 123.68])

        image.append(img)

    data = np.array(image)
    data = data.astype('float32')
    data /= 255

    return data


def train_data_generator(self, addresses, labels, batch_size):
    """Train data generator"""
    #Use first %80 of data for training.
    addresses = addresses[: int(0.8 * len(addresses))]
    labels = labels[: int(0.8 * len(labels))]
    total_data = len(addresses)
    while 1:
        for i in range(total_data / batch_size):
            batch_addresses = addresses[i * batch_size: (i + 1) * batch_size]
            batch_labels = labels[i * batch_size: (i + 1) * batch_size]

            data = self.prepare_input_data(batch_addresses)

            batch_labels = np_utils.to_categorical(batch_labels, self.nb_class)

            yield data, batch_labels

def val_data_generator(self, addresses, labels, batch_size):
    """Validation data generator"""
    #Use the last %20 of data for validation
    addresses = addresses[int(0.8 * len(addresses)):]
    labels = labels[int(0.8 * len(labels)):]
    total_data = len(addresses)
    image = []
    while 1:
        for i in range(total_data / batch_size):
            batch_addresses = addresses[i * batch_size: (i + 1) * batch_size]
            batch_labels = labels[i * batch_size: (i + 1) * batch_size]

            data = self.prepare_input_data(batch_addresses)

            batch_labels = np_utils.to_categorical(batch_labels, self.nb_class)

            yield data, batch_labels

def train(self, label_interested_in):
    """Trains the model"""
    #Read training data from json file, and get addresses and labels
    addresses, labels = self.create_address_and_label(label_interested_in)
    batch_size = 50
    train_batch_size = 40
    val_batch_size = 10
    steps = int(len(addresses) / batch_size) + 1
    print(len(addresses), steps)
    #Perform training
    history = self.custom_vgg_model.fit_generator(
        self.train_data_generator(addresses, labels, train_batch_size),
        steps_per_epoch=steps, epochs=self.number_of_epochs,
        verbose=1, validation_data=self.val_data_generator(addresses, labels, val_batch_size),
        validation_steps=steps, initial_epoch=0)

Why am I seeing such high memory usage? Is it because the way generators work in keras? I read that generators prepare batches beforehand to speedup the training process by running in parallel with the training. Or am I doing something wrong?

As a side question, since there isn't a batch_size argument in fit_generator(), am I correct in assuming that data gets loaded into the model based on generators and gradient updates are performed after each training and validation batch is loaded?

Upvotes: 2

Views: 3670

Answers (1)

Atta Jutt
Atta Jutt

Reputation: 490

Try workers=0

This will not invoke any multiprocessing which is intended to fill up the queue beforehand up to the max_queue_size argument with using k workers. What this does is; prepare a queue of generated data on CPU while training is ongoing on GPU so no time is lost and avoid bottlenecks.

For your need workers=0 will work

For deeper inquiry refer to keras fit_generator

Upvotes: 2

Related Questions