docstr: Compute loss from ``prediction`` and ``ground truth``,
  the computed loss map are weighted by ``weight_map``.
  
  if ``prediction`` is list of tensors, each element of the list
  will be compared against ``ground_truth` and the weighted by
  ``weight_map``.
  
  :param prediction: input will be reshaped into
      ``(batch_size, N_voxels, num_classes)``
  :param ground_truth: input will be reshaped into
      ``(batch_size, N_voxels)``
  :param weight_map: input will be reshaped into
      ``(batch_size, N_voxels)``
  :return:
docstr_tok: ["Compute","loss","from","``","prediction","``","and","``","ground","truth","``",",","the","computed","loss","map","are","weighted","by","``","weight_map","``",".","if","``","prediction","``","is","list","of","tensors",",","each","element","of","the","list","will","be","compared","against","``","ground_truth","`","and","the","weighted","by","``","weight_map","``",".",":","param","prediction",":","input","will","be","reshaped","into","``","(","batch_size",",","N_voxels",",","num_classes",")","``",":","param","ground_truth",":","input","will","be","reshaped","into","``","(","batch_size",",","N_voxels",")","``",":","param","weight_map",":","input","will","be","reshaped","into","``","(","batch_size",",","N_voxels",")","``",":","return",":"]
code: def layer_op(self, prediction, ground_truth=None, weight_map=None):
      """
      Compute loss from ``prediction`` and ``ground truth``,
      the computed loss map are weighted by ``weight_map``.
  
      if ``prediction`` is list of tensors, each element of the list
      will be compared against ``ground_truth` and the weighted by
      ``weight_map``.
  
      :param prediction: input will be reshaped into
          ``(batch_size, N_voxels, num_classes)``
      :param ground_truth: input will be reshaped into
          ``(batch_size, N_voxels)``
      :param weight_map: input will be reshaped into
          ``(batch_size, N_voxels)``
      :return:
      """
  
      with tf.device("/cpu:0"):
          batch_size = ground_truth.shape[0].value
          dir_size = 1
          if self._reshape:
              ground_truth = tf.reshape(ground_truth, [batch_size, -1])
              if weight_map is not None:
                  weight_map = tf.reshape(weight_map, [batch_size, -1])
          else:
              dir_size = ground_truth.shape[-1].value
              ground_truth = tf.reshape(
                  ground_truth, [batch_size, -1, dir_size]
              )
          if not isinstance(prediction, (list, tuple)):
              prediction = [prediction]
  
          data_loss = []
          for ind, pred in enumerate(prediction):
              # go through each scale
              def _batch_i_loss(*args):
                  # go through each image in a batch
                  if len(args[0]) == 2:
                      pred_b, ground_truth_b = args[0]
                      weight_map_b = None
                  else:
                      pred_b, ground_truth_b, weight_map_b = args[0]
                  pred_b = tf.reshape(pred_b, tf.shape(ground_truth_b))
                  # pred_b = tf.reshape(pred_b, [-1])
                  # pred_b = tf.Print(tf.cast(pred_b, tf.float32),
                  #                       [tf.shape(
                  #                           pred_b), tf.shape(
                  #                           ground_truth_b)],
                  #                       message='pred_b_shape')
  
                  loss_params = {
                      "prediction": pred_b,
                      "ground_truth": ground_truth_b,
                      "weight_map": weight_map_b,
                  }
                  if self._loss_func_params:
                      loss_params.update(self._loss_func_params)
  
                  return tf.to_float(
                      self._data_loss_func(**loss_params)
                  )
  
              if weight_map is not None:
                  elements = (pred, ground_truth, weight_map)
              else:
                  elements = (pred, ground_truth)
  
              loss_batch = tf.map_fn(
                  fn=_batch_i_loss,
                  elems=elements,
                  dtype=tf.float32,
                  parallel_iterations=1,
              )
              data_loss.append(tf.reduce_mean(loss_batch))
          return tf.reduce_mean(data_loss)
  
file_path: NiftyNet/niftynet/layer/loss_regression.py
imports: from __future__ import absolute_import, print_function, division
  import tensorflow
  from niftynet.engine.application_factory import LossRegressionFactory
  from niftynet.layer.base_layer import Layer
label_desc: Use image node_id 2 for calling the LossFunction obj's underlying member method code with example usage: obj.layer_op(prediction, ground_truth, weight_map) and returns: tf, tf