How to use a NN architecture that is too big for GPU?

Question

Initially posted in Stack Overflow.

I would like to implement a model which is actually 2 neural networks stacked together. However the size of these 2 architecture is too big to fit in GPU at the same time.

My idea was the following :

• Load the first model and run it for 1 batch
• Unload the first model from GPU, load the second model
• Run the second model from the output of the first model
• Unload the second model from GPU
• Repeat for every batch

I actually don't need to train the first model, since it's pre-trained. But I need to train second model.

Is it possible to do something like this ? Is my approach correct ? What are the pitfalls I should be aware of ? What about performance ?

Edit

I already tried the idea of computing the output of the first model for the whole dataset at first, and then use it as input for the second dataset. However the output of the first dataset is really big, and I don't have the available space for storing the whole pre-processed dataset. That is why I wanted to do it each batch.

Edit 2

After the very nice answer from @Gal Avineri, just one more precision : I would like to implements my architecture using only one GPU.

Answer 1

I suggest a method similar to what @ignatius offered.

Since you don't need to train the first model and only the second one you could do the following:

1. Use the first model over the entire dataset and save the extracted results in your memory.
2. Train the second model using the results from the previous step.

This way you will only have to load one model at a given time.
In addition this will make the training of the second model to be faster.

Edit 1

(the op has specified he does not have enough memory for the inference results)
In this case I can suggest to parallelize between the inference of the first model and the training of the second model.
This could be achieved through prefatching.

Let's denote the original dataset as data1 and the corresponding inference results from the first model as data2.
Let's also denote the first model as M1 and the second model as M2.
Thus you can do the following:

1. load M1 to cpu, and M2 to gpu.
2. Draw a batch from data1 and use M1 to prepare a batch of data2 for M2.
3. Use the gpu to train M2 on the batch received from the previous step
4. while step 3 is being executed, prepare the next batch of data2 by executing step 2.

In this way, while M2 is being trained on a batch, the next batch is being prepared. This is called prefetching the next batch.
This will parallelize the training of M2 and the making of data2.

The method i described above has a very simple way to implement using the tensorflow "Data" api.

Here is a code example:

import numpy as np
from tensorflow.keras import backend as K
from tensorflow.keras.layers import Input, Dense
from tensorflow.keras.models import Model
from tensorflow.data import Dataset

with K.device('cpu0'):
    input = Input((10,))
    x = Dense(100)(input)
    m1 = Model(input_shape=input, outputs=x)

with K.device('gpu0'):
    input = Input((100,))
    x = Dense(5)(input)
    m2 = Model(input_shape=input, outputs=x)


features = np.random.rand(10, 2000)
labels = np.random.randint(2, size=2000)
dataset = Dataset.from_tensor_slices((features, labels))


def preprocess_sample(features, label):
    inference = m2.predict(features, batch_size=1)
    return inference, label


dataset = dataset.shuffle(2000).repeat().map(preprocess_sample).batch(32)
dataset = dataset.prefetch(1)

m2.fit(dataset, epochs=100, steps_per_epoch=2000//32)

Answer 2

I've seen this problem before with big CNN architectures. If the first network isn't going to change, then do a prediction with that network on all the data and use the output as the training data for the second network. You can save this as a separate dataset and reuse it.

Edit- I just noticed Gal Avineri posted the same answer.