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 ?


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.

  • $\begingroup$ I'ver never come across with something similar, but, have you though about loading the first model in the RAM (not gpu-ram) and perform the inference in the CPU for the batch?, then you can insert the output of the first model into the second model, which is trained on the GPU. It sounds quite inefficient, but unloading and loading two entire networks in each batch seems so too... $\endgroup$
    – ignatius
    Nov 15, 2018 at 9:49
  • $\begingroup$ I think you'd relieved out of the stress if you hire 1 virtual server with GPU from AWS EC2, Digital Ocean or Google Cloud, and run them together with adding a simple code that merges their network (even SFTP would manage it) to enabling to transmission of the outputs. Of course that would take money, but I do not think it is efficient in your way but it seems still possible to implement though. $\endgroup$
    – Ugur MULUK
    Nov 15, 2018 at 11:02

2 Answers 2


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)
  • $\begingroup$ This is what I wanted to do at first. However (my bad I didn't precise the question), the output of the first model is quite big, and I don't have the space necessary for storing such a big file for the whole dataset... $\endgroup$
    – Astariul
    Nov 15, 2018 at 23:23
  • 1
    $\begingroup$ I've updated my answer according to your edit :) $\endgroup$ Nov 16, 2018 at 14:32
  • $\begingroup$ Very helpful edit, code is very nice. Thank you. I edited the question again ^^ $\endgroup$
    – Astariul
    Nov 18, 2018 at 23:40
  • 1
    $\begingroup$ I'm glad you find my idea helpful to you :) I've seen your update, and it seems to me the code i have proposed also fits your new request. In the proposed code, m1 is loaded to the cpu and m2, which should be trained, is loaded into a single gpu (which is your first gpu, also notated as gpu0) :). Can you elaborate why the currently proposed code doesn't fit your requirements? :) $\endgroup$ Nov 19, 2018 at 8:26
  • 1
    $\begingroup$ In that case you have only one resource and you must at any given moment run either m1 or m2. If the time required to load a model into the gpu is much longer than executing a forward & backward routine on a batch, than the optimal solution would be to switch the model as few times as possible. Therefore i would suggest to load m1, prepare as many batches as your memory can save, than load m2 to train on each of them once, and repeat this process $\endgroup$ Nov 20, 2018 at 11:27

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.


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.