How do I run an iterative 2D convolution for each slice of a tensor?

user_16434104 0
Viewed 247

I'm working on a machine learning project with convolutional neural networks using TF/Keras in Python, and my goal is to split up an image up into patches, run a convolution on each one separately, and then put it back together.

What I can't figure out how to do is run a convolution for each slice of a 3D array.

For example, if I have a tensor of size (500,100,100) I want to do a separate convolution for all 500 slices of size (100 x 100). I'm implementing this within a custom Keras layer and want these to be trainable weights I've tried a few different things:

  1. Using map.fn() to run a convolution for each slice of the array
    • This doesn't seem to attach weights to each layer separately.
  2. Using the DepthwiseConv2D layer:
    • This works well for the first call of the layer, but fails when I call the layer the second time with more filters because it wants to perform the depthwise convolution on each of the previous filtered layers
    • This, of course isn't what I want because I want one convolution for each of the previous sets of filters from the previous layer.

Any ideas are appreciated, as I'm truly stuck here. Thank you!

keras tensorflow deep-learning neural-network python
1 Answers

If you have a tensor with shape (500,100,100) and want to feed some subset of this tensor, to separate conv2d layers at the same time, you may do this by defining conv2d layers in the same level. You should first define Lambda layers to split input, then feed their output to Conv2D layers, then concatenate them.

Let's take a tensor with shape (100,28,28,1) as an example, that we want to split it into 2 subset tensor and apply conv2d layers on each subset separately:

import tensorflow as tf
from tensorflow.keras.layers import Dense, Flatten, Conv2D, Input, concatenate, Lambda
from tensorflow.keras.models import Model

# define a sample dataset
x = tf.random.uniform((100, 28, 28, 1))
y = tf.random.uniform((100, 1), dtype=tf.int32, minval=0, maxval=9)
ds = tf.data.Dataset.from_tensor_slices((x, y))
ds = ds.batch(16)

def create_nn_model():
    input = Input(shape=(28,28,1))
    b1 = Lambda(lambda a: a[:,:14,:,:], name="first_slice")(input)
    b2 = Lambda(lambda a: a[:,14:,:,:], name="second_slice")(input)
    d1 = Conv2D(64, 2, padding='same', activation='relu', name="conv1_first_slice")(b1)
    d2 = Conv2D(64, 2, padding='same', activation='relu', name="conv2_second_slice")(b2)
    x =  concatenate([d1,d2], axis=1)
    x = Flatten()(x)
    x = Dense(64, activation='relu')(x)
    out = Dense(10, activation='softmax')(x)
    model = Model(input, out)
    model.compile(loss='sparse_categorical_crossentropy', optimizer='adam', metrics=['accuracy'])
    return model

model = create_nn_model()
tf.keras.utils.plot_model(model, show_shapes=True)

Here is the plotted model architecture:

plotted_model

user_24232780
Related