How to define a weighted loss function in TensorFlow?
<p>I have a training dataset of train_data and train_labels which is train_data_node and train_labels_node in the graph of tensorflow.
As you know, I can use the loss function of tensorflow as bellows:</p>

<pre><code>logits = model(train_data_node)
loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(
    logits,train_labels_node))
</code></pre>

<p>However, this loss function processes all the training data equally.
But in our situation, we want to process the data discriminately.
For example, we have a csv file corresponding to the training data to indicate the train data is original or augmented.
Then we want to define a custom loss function which makes the loss of original data play more important role and the loss of augmented data play less important role, such as:</p>

<pre><code>loss_no_aug = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(noAugLogits, noAugLabels))
loss_aug = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(augLogits, augLabels))
loss = loss_no_aug * PENALTY_COEFFICIENT + loss_aug
</code></pre>

<p>I have defined a loss function as bellow, but it didn't work:</p>

<pre><code>def calLoss(logits, labels, augs):
  noAugLogits = []
  noAugLabels = []
  augLogits = []
  augLabels = []
  tf.get_collection()
  for i in range(augs.shape[0]):
    if augs[i] == 1:
      noAugLogits.append(logits[i])
      noAugLabels.append(labels[i])
    else:
      augLogits.append(logits[i])
      augLabels.append(labels[i])
  noAugLogits = tf.convert_to_tensor(noAugLogits)
  noAugLabels = tf.convert_to_tensor(noAugLabels)
  augLogits = tf.convert_to_tensor(augLogits)
  augLabels = tf.convert_to_tensor(augLabels)
  return tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(
      noAugLogits, noAugLabels)) * PENALTY_COEFFICIENT + \
      tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(augLogits, augLabels))
</code></pre>

<p>I think we should write the loss function using tensor operations, however, I am not familiar with them. So could anyone give me some advice on how to define the loss function.</p>

<p>Thank you for your kind answers or suggestions.</p>