Creating parallel keras layers

Question

I am new to Keras and ML and I want to create a NN that can seperate a bitmap-like image into its visual components.

My approach is to feed a two dimensional image (lets say 8 by 8 pixels) into a NN, which then outputs a three dimensional matrix (e.g 3 by 8 by 8). Every element of the first dimension represents an image of a component (see this illustration).

As a first version I created a Keras sequential model with 1 flatten layer and 2 dense layers, 192 units each (3x8x8=192). After around 5000 training images, the model performance was still mediocre.

Question: Is there a way of splitting the output of the last dense layer into three seperate images that will then be processed in a parallel manner by further dense layers? Is there a better approach to the problem I am facing?

Answer 1

So, yes and no.

First, there's not a layer that does this in the standard Keras API. It might be possible to write a custom layer to do it, but I'm not comfortable enough doing that to guide you.

What you COULD do is create three different layers, each of which accepts the whole thing as input, and let each of them figure out what parts they're responsible for. More on that in a minute.

Second, to do it requires you to use the Functional API. The Sequential API is just that - sequential. But functional is pretty easy once you understand what's going on.

Ok. Now for the above approach.

import tensorflow as tf
from tensorflow.keras.layers import Input, Dense, Flatten, Concatenate
from tobias_code import get_my_data, compiler_args

data = get_my_data() # obviously, this is a stand-in for however you get your data.

input_layer = Input(data.shape[1:])
hidden = Flatten()(input_layer)

This is where the Functional API differs from the Sequential API. Because the Sequential API assumes each layer attaches to the previous layer, you don't have to manually do so. The Functional API allows you to connect layers however you'd like! But to allow for that, it requires you to manually connect each layer to the previous. I typically reuse 'hidden' to chain my layers together, deviating only for the input, output, and any unusually connected layers.

hidden = Dense(192, activation='relu')(hidden)
main_output = Dense(192, activation='relu')(hidden)

Here's where yours would end. You'd call model = tf.keras.Model(input_layer, main_output), do whatever compiling you had to do, and begin training. But we're going to continue on. (You probably wouldn't use relu for your output layer, but eh...)

# I'm going to build each individual parallel set of layers separately
branch_a = Dense(96, activation='relu')(main_output)
branch_a = Dense(48, activation='relu')(branch_a)
output_a = Dense(24)(branch_a) # This will be one of my outputs, so I want a linear activation

branch_b = Dense(96, activation='relu')(main_output) # note that it is main_output again
branch_b = Dense(48, activation='relu')(branch_b)
output_b = Dense(24)(branch_b)

branch_c = Dense(96, activation='relu')(main_output) # and again. 3 layers are all sharing that output.
branch_c = Dense(48, activation='relu')(branch_c)
output_c = Dense(24)(branch_c)

Now, the Functional API can handle multiple outputs. You can define separate loss functions for each output and separate metrics for each output. But I haven't done a lot with that, so rather than tell you incorrectly how to do it, we're going to concatenate your outputs into one, which is probably what your output looked like in the sequential model.

final_output = Concatenate()([output_a, output_b, output_c])
model = tf.keras.Model(input_layer, final_output)
model.compile(**compiler_args)
model.summary()

Now, I haven't run it, and I make no claims as to its efficacy; I pretty much typed it here and am hoping for the best. But it should get you where you want to be, and hopefully if it's not exactly what you want it will set you on the right path.