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:


enter image description here

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: enter image description here

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)

Loss graph after training stage will looks like this: enter image description here

But if I change offset to 1e-4 I will get this loss graph: enter image description here 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.


1 Answer 1


One possible issue that you are switching between tf.math.atan2 and np.arctan2. Only use tf.math.atan2 because that function is designed to work with Keras.


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.