Search code examples
tensorflowloss-function

Tensorflow, proper way to define a loss function with element-wise operations


I have a batch of n training examples. Each training example yields me with a vector a (of length k) and a vector c (also of length k).

In tensorflow my output is p.

I would like to define the following loss function:

\sum_(i = 0)^(i=n) c_i ( (a_i - p)^2 )

I have read about various tf operations available to me but cannot seem to find out a proper implementation for this. I tried to use tf.tile to replicate p into a k-length tensor and do a tf.subtract, but this seems excessive.

Any help would be much appreciated.


Solution

  • What is the shape of p? TensorFlow natively supports subtraction of tensors with different shapes:

    import tensorflow as tf
    
    a = tf.Variable([[1, 2, 3], [2, 3, 4], [3, 4, 5]])
    
    x_1 = a - tf.Variable([1, 1, 1])
    x_2 = a - tf.Variable([1])
    x_3 = a - tf.Variable(1)
    
    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        print(sess.run(x_1))  # [[0, 1, 2], [1, 2, 3] [2, 3, 4]]
        print(sess.run(x_2))  # [[0, 1, 2], [1, 2, 3] [2, 3, 4]]
        print(sess.run(x_3))  # [[0, 1, 2], [1, 2, 3] [2, 3, 4]]
    

    You can define your loss function as follows:

    a = tf.Variable([[1, 2, 3], [2, 3, 4], [3, 4, 5]])
    c = tf.Variable([[4, 5, 6], [5, 6, 7], [6, 7, 8]])
    
    p = tf.Variable(1)
    
    loss_1 = tf.reduce_mean(tf.reduce_sum(tf.multiply(c, tf.pow(tf.subtract(a, p), 2)), axis=1))  # 112
    loss_2 = tf.reduce_mean(tf.reduce_sum(c * (a-p)**2, axis=1))  # 112