Splitting a neural network in 2 microservices

Question

I have a neural network, that is already trained locally that can detect objects in the scene.

But I have to split the neural network into 2 parts, let's say it has 16 layers, and I want to have one microservice handle the first 8 layers and give the output of the 8th layer to the next microservice and it takes the data from there and proceeds with the 9th layer (1st layer in the 2nd microservice).

Sending an image to the 1st microservice will give the result from the 2nd microservice, is this feasible using TensorFlow?

Answer 1

Yes, you can split the model into two parts after training. Not sure, what would be the advantage, but it is possible. High level steps :

1. Define and train the model
2. Save trained model
3. Create two copies of model (Layers 1 to 8 and new input layer + 9 to 16)
4. Transfer weights to these copies of model

After this, service 1 can run first few layers and provide the output to service 2.

Example code to split the model (Code uses Keras as wrapper over TensorFlow):

import keras
from keras.models import Model, load_model
from keras.layers import Input, Dense
from keras.optimizers import RMSprop
import numpy as np


# Create original model and save it
inputs = Input((1,))
dense_1 = Dense(10, activation='relu')(inputs)
dense_2 = Dense(10, activation='relu')(dense_1)
dense_3 = Dense(10, activation='relu')(dense_2)
outputs = Dense(10)(dense_3)

model = Model(inputs=inputs, outputs=outputs)
model.compile(optimizer=RMSprop(), loss='mse')
model.save('test.h5')


# Load the model and make modifications to it
loaded_model = load_model('test.h5')
loaded_model.layers.pop()
loaded_model.layers.pop()


# Create your new model with the two layers removed and transfer weights
new_model = Model(inputs=inputs, outputs=dense_1)
new_model.compile(optimizer=RMSprop(), loss='mse')
new_model.set_weights(loaded_model.get_weights())

new_model.summary()
new_model.save('test_complete.h5')

Source : https://github.com/keras-team/keras/issues/8772