custom_layer.py

import tensorflow as tf
from tensorflow import keras
import numpy as np
from tensorflow.keras.models import Sequential
from tensorflow.keras.optimizers import RMSprop
from tensorflow.keras.optimizers import SGD
from tensorflow.keras.layers import Flatten

ig = 0.0
im = 0.5
iw = 0.88
isg = 0.2
ism = 0.2
isw = 0.2
eps=1e-7
pi=3.14
epoch_count = 1000
learning_rate = 0.001
sequence_len = 100

def w_initialization(dimension, activationrange = 2):
    if type(dimension) is list and len(dimension) > 1:
        input_dim = dimension[0]
        output_dim = dimension[1]
    elif type(dimension) is list:
        input_dim = 0
        output_dim = dimension[0]
    else:
        input_dim = 0
        output_dim = dimension
    return np.random.normal(loc=iw, scale=isw, size=dimension).astype('float32')

class NMU(keras.layers.Layer):
    def __init__(self, output_dim, input_dim):
        super(NMU, self).__init__()
        self.weight = tf.Variable(w_initialization([input_dim, output_dim]), dtype="float32", name="weight", trainable=True)
    def call(self, inputs):
        W = tf.minimum(tf.maximum(self.weight , 0), 1)
        output = (inputs * W)+1-W
        return tf.reduce_prod(output, axis=1)

X_TRAIN = []
X_TEST = []

Y_TRAIN = []
Y_TEST = []

#DATASET
for x in range(1, 10000):
    a1 = np.random.randint(1000)
    a2 = np.random.randint(1000)
    X_TRAIN.append(np.array([np.float32(a1), np.float32(a2)]))
    Y_TRAIN.append(np.array([np.float32(a2*a1), np.float32(a2*a1)]))

for x in range(1, 10000):
    a1 = np.random.randint(1000)
    a2 = np.random.randint(1000)
    X_TEST.append(np.array([np.float32(a1),np.float32(a2)]))
    Y_TEST.append(np.array([np.float32(a2*a1), np.float32(a2*a1)]))

X_TRAIN = np.array(X_TRAIN)
X_TEST = np.array(X_TEST)
Y_TEST = np.array(Y_TEST)
Y_TRAIN = np.array(Y_TRAIN)

#MODEL
model = Sequential()
model.add(NMU(2,2))
model.compile(loss='mse', optimizer=RMSprop(lr=learning_rate))
model.fit(
    validation_data=(X_TEST, Y_TEST),
    x=X_TRAIN,
    y=Y_TRAIN,
    epochs=epoch_count
)