Creating Custom Layer As Stack Of Individual Neurons Tensorflow

So, I'm trying to create a custom layer in TensorFlow 2.4.1, using a function for a neuron I defined. # NOTE: this is not the actual neuron I want to use, # it's just a simple exam

Solution 1:

So, person who asked this question here, I have found a way to do it, by dynamically creating variables and operations. First, let's re-define the neuron to use tensorflow operations:

def neuron(x, W, b):
    return tf.add(tf.matmul(W, x), b)

Then, let's create the layer (this uses the blueprint layed out in the question):

classLayer(tf.keras.layers.Layer):
    def__init__(self, n_units=5):
        super(Layer, self).__init__()

        self.n_units = n_units

    defbuild(self, input_shape):
        for i inrange(self.n_units):
            exec(f'self.kernel_{i} = self.add_weight("kernel_{i}", shape=[1, int(input_shape[0])])')
            exec(f'self.bias_{i} = self.add_weight("bias_{i}", shape=[1, 1])')

    defcall(self, inputs):
        for i inrange(self.n_units):
            exec(f'out_{i} = neuron(inputs, self.kernel_{i}, self.bias_{i})')
        returneval(f'tf.concat([{", ".join([ f"out_{i}"for i inrange(self.n_units) ])}], axis=0)')

As you can see, we're using exec and eval to dynamically create variables and perform operations. That's it! We can perform a few checks to see if TensorFlow could use this:

# Check to see if it outputs the correct thing
layer = Layer(5) # With 5 neurons, it should return a (5, 6)print(layer(tf.zeros([10, 6])))

# Check to see if it has the right trainable parametersprint(layer.trainable_variables)

# Check to see if TensorFlow can find the gradients
layer = Layer(5)
x = tf.ones([10, 6])
with tf.GradientTape() as tape:
    z = layer(x)
print(f"Parameter: {layer.trainable_variables[2]}")
print(f"Gradient:  {tape.gradient(z, layer.trainable_variables[2])}")

This solution works, but it's not very elegant... I wonder if there's a better way to do it, some magical TF method that can map the neuron to create a layer, I'm too inexperienced to know for the moment. So, please answer if you have a (better) answer, I'll be happy to accept it :)