title: How to define a weighted loss function in TensorFlow?
body: <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>