I'm trying to add arctan2 function to the end of Keras model, but it looks like it is not getting any near even local minimum. Here is my ridiculous but minimal working code with native Keras Add() function instead of arctan2 function:
import numpy as np
import matplotlib.pyplot as plt
import scipy.signal as ss
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input, Conv1D, Add
import tensorflow as tf
kernel_size = 64
epochs = 1000
def atan2(tensors):
Q = tensors[0]
I = tensors[1]
return tf.math.atan2(Q, I)
def atan2_output_shape(input_shapes):
return tuple(input_shapes[0])
atan2_layer = Lambda(atan2, output_shape=atan2_output_shape)
## Data generation for training
x_train = np.random.randn(1024, 512)
t = np.linspace(0, x_train.shape[1], x_train.shape[1], endpoint=False)
sine = np.sin(2*np.pi*t/32)
cosine = np.cos(2*np.pi*t/32)
x_I = np.multiply(x_train, cosine)
x_Q = np.multiply(x_train, sine)
b_I = ss.tukey(kernel_size)
b_Q = ss.tukey(kernel_size)
x_I_filt = np.array([np.convolve(b_I, x_I_i, mode='valid') for x_I_i in x_I])
x_Q_filt = np.array([np.convolve(b_Q, x_Q_i, mode='valid') for x_Q_i in x_Q])
y_train = x_Q_filt + x_I_filt
# y_train = x_Q_filt * x_I_filt
# y_train = np.arctan2(x_Q_filt, x_I_filt)
x_I = np.expand_dims(x_I, axis=2)
x_Q = np.expand_dims(x_Q, axis=2)
y_train = np.expand_dims(y_train, axis=2)
## Keras model
input_I = Input(shape=(x_I.shape[1], 1))
input_Q = Input(shape=(x_Q.shape[1], 1))
conv_I_1D = Conv1D(filters=1, kernel_size=kernel_size, activation=None, padding='valid', use_bias=False)(input_I)
conv_Q_1D = Conv1D(filters=1, kernel_size=kernel_size, activation=None, padding='valid', use_bias=False)(input_Q)
out_I_Q = Add()([conv_I_1D, conv_Q_1D])
# out_I_Q = Multiply()([conv_I_1D, conv_Q_1D])
# out_I_Q = atan2_layer([conv_Q_1D, conv_I_1D])
model_1D = Model([input_I, input_Q], out_I_Q)
model_1D.compile(optimizer='sgd', loss='mean_squared_error')
history_1D = model_1D.fit([x_I, x_Q], y_train, epochs=epochs, verbose=0)
and after 100 epochs I'm getting almost perfect initial filter kernel:
plt.semilogy(history_1D.history['loss'])
plt.ylabel('Loss')
plt.xlabel('Epoch')
plt.show()
It also works well with Multiply() function. But if I replace y_train to y_train = np.arctan2(x_Q_filt, x_I_filt) and out_I_Q to out_I_Q = atan2_layer([conv_Q_1D, conv_I_1D]) I will get this sad loss graph:
I even initialize weights as they should be but with a little bit of an offset before runningmodel_1D.fit(...). b_I and b_Q are the same arrays.
offset = 1e-5
array_for_I_weights = np.array(model_1D.layers[2].get_weights())
array_for_I_weights[0,:,0,0] = list(b_I+offset)
model_1D.layers[2].set_weights(array_for_I_weights)
model_1D.layers[3].set_weights(array_for_I_weights)
Loss graph after training stage will looks like this:
But if I change offset to 1e-4 I will get this loss graph:
It only getting worse epoch after epoch. What is wrong with arctan2 function in case of placing it inside of Keras model? Why model is jumping away from good minimum? Maybe I should use another custom loss/metric function?
Changing optimizer type won't help at all. I'm using tensorflow 2.0 and Keras 2.2.5.
