Is it possible to build a custom loss function of chamfer distance using Keras - keras

Question

In my assignment about the point cloud,I need to use a keras custom loss function of chamfer distance and apply it to the autoencoder.But I find it is hard to implement the function.
I have tried to write in this way.I hope someone could help me.
def chamfer_loss_func(y_true,y_pred):
'''
Calculate the chamfer distance,use euclidean metric
:param y_true:
:param y_pred:
:return:
'''
batch_size = 32
y_true = K.reshape(y_true,[batch_size,2048,3])
y_pred = K.reshape(y_pred, [batch_size, 2048, 3])
num_t = K.int_shape(y_true)[1]
num_p = K.int_shape(y_pred)[1]
dists_mat = K.zeros(shape=[num_t, num_p])
sum = 0.0
for bi in range(batch_size):
for i in range(num_t):
pnt_t = y_true[bi][i]
for j in range(num_p):
if (i <= j):
pnt_p = y_pred[bi][j]
dists_mat[i][j] = K.eval(K.sum(K.sqrt(tf.subtract(pnt_t, pnt_p))))
else:
dists_mat[i][j] = dists_mat[j][i]
dist_t_to_p = K.mean(K.min(dists_mat, axis=0))
dist_p_to_t = K.mean(K.min(dists_mat, axis=1))
sum += K.max([dist_p_to_t, dist_t_to_p])
return sum / batch_size
And the network structure is listed here:
def get_model(num_pnts):
input_tensor = Input(shape=[num_pnts, 3])
input_cov_tensor = Input(shape=[num_pnts, 9]) # add the local covariance matrix
# the encoder network
concat_1 = concatenate([input_tensor, input_cov_tensor]) # concatenate the two input tensor
dense_1 = Dense(32,activation='relu')(concat_1)
dense_2 = Dense(64,activation='relu')(dense_1)
dense_3 = Dense(128,activation='relu')(dense_2)
encoded = MaxPooling1D(pool_size=num_pnts)(dense_3)
# the decoder network
# use 3 fully connected layers
dense_3 = Dense(128,activation='relu')(encoded)
dense_4 = Dense(128,activation='relu')(dense_3)
dense_5 = Dense(num_pnts*3,activation='linear')(dense_4)
decoded = Reshape(target_shape=[num_pnts,3])(dense_5)
#the autoencoder
autoencoder = Model(inputs=[input_tensor,input_cov_tensor],outputs=decoded)
return autoencoder
And the place where uses the loss function:
model = get_model(2048)
model.compile(optimizer='adam',loss=chamfer_loss_func,)

Answer 1

Can you give me the link of chamfer distance you reference to?
K.sum(K.sqrt(tf.subtract(pnt_t, pnt_p))) looks strange to me. To calculate Euclidean distance, the sqrt should be replaced with square.
And it is not recommended to use for loop in tensorflow, so I have reimplemented it:
import numpy as np
import keras.backend as K
import tensorflow as tf
from keras.layers import Input, Dense, MaxPooling1D, Reshape, concatenate
from keras.models import Model
def dists_mat_calculater(pnts_t, pnts_p):
"""
arguments:
pnts_t : from y_true[bi], shape: (num_t, 3)
pnts_p : from y_pred[bi], shape: (num_p, 3)
return:
dists_mat: shape: (num_t, num_p)
"""
num_t = K.int_shape(pnts_t)[0]
num_p = K.int_shape(pnts_p)[0]
pnts_t = tf.reshape(tf.tile(tf.expand_dims(pnts_t, 1), [1, 1, num_p]), [-1, 3])
pnts_p = tf.tile(pnts_p, [num_t, 1])
dists_mat = K.sum(K.square(tf.subtract(pnts_t, pnts_p)), axis=1)
dists_mat = tf.reshape(dists_mat, [num_t, num_p])
dists_mat_upper = tf.matrix_band_part(dists_mat, 0, -1)
dists_mat_symm = dists_mat_upper + tf.transpose(dists_mat_upper)
dists_mat_symm = tf.matrix_set_diag(dists_mat_symm, tf.diag_part(dists_mat))
return dists_mat_symm
def dist_calculator(pnts):
"""
arguments:
pnts_t : from y_true[bi], shape: (num_t, 3)
pnts_p : from y_pred[bi], shape: (num_p, 3)
return:
dist: shape: (1, )
"""
pnts_t, pnts_p = pnts
dists_mat = dists_mat_calculater(pnts_t, pnts_p)
dist_t_to_p = K.mean(K.min(dists_mat, axis=0)) #shape: (1,)
dist_p_to_t = K.mean(K.min(dists_mat, axis=1)) #shape: (1,)
dist = K.max([dist_p_to_t, dist_t_to_p]) #shape: (1,)
return dist
def chamfer_loss_func_tf(y_true,y_pred):
'''
Calculate the chamfer distance,use euclidean metric
:param y_true:
:param y_pred:
:return:
'''
y_true = K.reshape(y_true,[-1, num_pnts, 3])
y_pred = K.reshape(y_pred, [-1, num_pnts, 3])
return K.mean(tf.map_fn(dist_calculator, elems=(y_true, y_pred), dtype=tf.float32))
dists_mat_calculater calculate the distance matrix, and the part of pairwise calculation is enlightened by Mask R-CNN - overlaps_graph.
I also implement a pure python version for validation purpose:
def chamfer_loss_python(y_true,y_pred):
'''
Calculate the chamfer distance,use euclidean metric
:param y_true:
:param y_pred:
:return:
'''
y_true = np.reshape(y_true,[-1,num_pnts,3])
y_pred = np.reshape(y_pred, [-1,num_pnts, 3])
batch_size = y_true.shape[0]
num_t = y_true.shape[1]
num_p = y_pred.shape[1]
dists_mat = np.zeros((num_t, num_p))
_sum = 0.0
loss_before_mean_py = []
for bi in range(batch_size):
for i in range(num_t):
pnt_t = y_true[bi][i]
for j in range(num_p):
pnt_p = y_pred[bi][j]
if (i <= j):
pnt_p = y_pred[bi][j]
dists_mat[i][j] = np.sum((pnt_t - pnt_p)**2)
else:
dists_mat[i][j] = dists_mat[j][i]
dist_t_to_p = np.mean(np.min(dists_mat, axis=0))
dist_p_to_t = np.mean(np.min(dists_mat, axis=1))
_sum += np.max([dist_p_to_t, dist_t_to_p])
loss_before_mean_py.append(np.max([dist_p_to_t, dist_t_to_p]))
return _sum / batch_size
Following is the testing script:
num_pnts = 8
np.random.seed(1)
Y_true = np.random.randn(32, num_pnts, 3).astype(np.float32)
Y_pred = np.random.randn(32, num_pnts, 3).astype(np.float32)
Y_true_ph = tf.placeholder(tf.float32, shape=(None, num_pnts, 3), name="Y_true_ph")
Y_pred_ph = tf.placeholder(tf.float32, shape=(None, num_pnts, 3), name="Y_pred_ph")
loss = chamfer_loss_func_tf(Y_true_ph, Y_pred_ph)
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
loss = sess.run(loss, feed_dict={
Y_true_ph: Y_true,
Y_pred_ph: Y_pred})
loss_py = chamfer_loss_python(Y_true,Y_pred)
print(loss)
print(loss_py)