def _build_loss(self): train_paf, train_paf_mask = super().get_train_paf() train_heatmap, train_heatmap_mask = super().get_train_heatmap() train_mask = super().get_train_mask() paf_losses = [] heatmap_losses = [] for paf in super().get_paf_tensors(): paf_loss = Loss.abs_loss(train_paf, paf, raw_tensor=True) * train_paf_mask # --- LOSS MASKING paf_loss = paf_loss * tf.expand_dims(train_mask, axis=-1) paf_loss = paf_loss * train_paf_mask if self._paf_weight is not None: abs_training_paf = tf.abs(train_paf) mask = tf.cast( tf.math.greater(abs_training_paf, EPS), dtype=tf.float32 ) weights_mask = mask * self._paf_weight + self.__IDENTITY paf_loss = paf_loss * weights_mask paf_losses.append( tf.reduce_mean(paf_loss) ) for heatmap in super().get_heatmap_tensors(): heatmap_loss = Loss.abs_loss(train_heatmap, heatmap, raw_tensor=True) # --- LOSS MASKING heatmap_loss = heatmap_loss * tf.expand_dims(train_mask, axis=-1) heatmap_loss = heatmap_loss * train_heatmap_mask if self._heatmap_weight is not None: # Create mask for scaling loss # Add 1.0 for saving values that are equal to 0 (approximately equal to 0) weight_mask = train_heatmap * self._heatmap_weight + self.__IDENTITY heatmap_loss = heatmap_loss * weight_mask heatmap_losses.append( tf.reduce_mean(heatmap_loss) ) self._paf_loss = tf.reduce_sum(paf_losses) self._heatmap_loss = tf.reduce_sum(heatmap_losses) loss = self._heatmap_loss * self._heatmap_scale + \ self._paf_loss * self._paf_scale # For Tensorboard super().track_loss(self._paf_loss, PETrainer.PAF_LOSS) super().track_loss(self._heatmap_loss, PETrainer.HEATMAP_LOSS) return loss
def _build_loss(self): logits = super().get_logits() labels = super().get_labels() num_positives = None if self._normalize_by_positives: positives = tf.cast(tf.not_equal(labels, 0), tf.float32) # [BATCH_SIZE, ...] positives_dim_n = len(positives.get_shape()) axis = list(range(1, positives_dim_n)) num_positives = tf.reduce_sum( positives, axis=axis) # [BATCH_SIZE, N_POSITIVES] focal_loss = Loss.focal_binary_loss( logits=logits, labels=labels, num_positives=num_positives, focal_gamma=self._focal_gamma, label_smoothing=self._smoothing_labels) if not self._normalize_by_positives: focal_loss = focal_loss / float(super().get_batch_size()) super().track_loss(focal_loss, FocalBinaryTrainer.FOCAL_LOSS) return focal_loss
def _build_loss(self): logits = super().get_logits() labels = super().get_labels() ce_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=labels, logits=logits) num_classes = super().get_num_classes() logits_shape = tf.shape(logits) batch_size = logits_shape[0] positives = tf.not_equal(labels, 0) # [BATCH_SIZE, ...] positives_dim_n = len(positives.get_shape()) axis = list(range(1, positives_dim_n)) num_positives = tf.reduce_sum(positives, axis=axis) # [BATCH_SIZE, N_POSITIVES] maki_loss = Loss.maki_loss( flattened_logits=tf.reshape(logits, shape=[batch_size, -1, num_classes]), flattened_labels=tf.reshape(labels, shape=[batch_size, -1]), num_classes=num_classes, num_positives=num_positives, maki_gamma=self._maki_gamma, ce_loss=ce_loss) super().track_loss(maki_loss, MakiLossTrainer.MAKI_LOSS) return maki_loss
def _build_loss(self): labels = super().get_labels() logits = super().get_logits() # p - predicted probability # g - ground truth label p = tf.nn.sigmoid(logits) g = labels dice_loss = Loss.dice_loss( p=p, g=g, eps=self._eps, axes=self._axes ) super().track_loss(dice_loss, DiceTrainer.DICE_LOSS) return dice_loss
def _build_loss(self): logits = super().get_logits() labels = super().get_labels() num_classes = super().get_num_classes() positives = tf.not_equal(labels, 0) # [BATCH_SIZE, ...] positives_dim_n = len(positives.get_shape()) axis = list(range(1, positives_dim_n)) num_positives = tf.reduce_sum(positives, axis=axis) # [BATCH_SIZE, N_POSITIVES] focal_loss = Loss.focal_loss(logits=logits, labels=labels, num_classes=num_classes, num_positives=num_positives, focal_gamma=self._focal_gamma, raw_tensor=True) weights = super().get_weight_map() focal_loss = tf.reduce_sum(focal_loss * weights) super().track_loss(focal_loss, WeightedFocalTrainer.FOCAL_LOSS) return focal_loss
def _build_local_loss(self, prediction, label): abs_loss = Loss.abs_loss(label, prediction, raw_tensor=True) final_loss = tf.reduce_mean(abs_loss) return final_loss
def _build_loss(self): ce_loss = tf.nn.sparse_softmax_cross_entropy_with_logits( labels=super().get_labels(), logits=super().get_logits()) qce_loss = Loss.quadratic_ce_loss(ce_loss=ce_loss) super().track_loss(qce_loss, QCETrainer.QCE_LOSS) return qce_loss
def _build_loss(self): train_paf, train_paf_mask = super().get_train_paf() train_heatmap, train_heatmap_mask = super().get_train_heatmap() train_mask = super().get_train_mask() paf_losses = [] heatmap_losses = [] # --- PAF LOSS for paf in super().get_paf_tensors(): # Division by 2.0 makes it similar to tf.nn.l2_loss paf_loss = Loss.mse_loss(train_paf, paf, raw_tensor=True) / 2.0 # --- LOSS MASKING paf_loss = paf_loss * train_mask paf_loss = paf_loss * train_paf_mask if self._paf_weight is not None: abs_training_paf = tf.abs(train_paf) mask = tf.cast(tf.math.greater(abs_training_paf, EPS), dtype=tf.float32) weights_mask = mask * self._paf_weight + self.__IDENTITY paf_loss = paf_loss * weights_mask paf_losses.append(tf.reduce_sum(paf_loss)) # --- HEATMAP LOSS for heatmap in super().get_heatmap_tensors(): # We need to expand dims first, because tf.keras.losses.binary_crossentropy averages # the loss along the last dimension: [BS, H, W, C] -> [BS, H, W] # [BS, H, W, C, 1] train_heatmap_expanded = tf.expand_dims(train_heatmap, axis=-1) heatmap_expanded = tf.expand_dims(heatmap, axis=-1) heatmap_loss = tf.keras.losses.binary_crossentropy( y_true=train_heatmap_expanded, y_pred=heatmap_expanded, label_smoothing=self._label_smoothing) / 2.0 # --- LOSS MASKING heatmap_loss = heatmap_loss * train_mask heatmap_loss = heatmap_loss * train_heatmap_mask # heatmap_loss - [BS, H, W, C] if self._is_nullify_absent_labels: # [bs, 1, 1, C] label_sum = tf.reduce_sum(train_heatmap * train_mask, axis=[1, 2], keepdims=True) scale_factor = label_sum / (label_sum + 1e-5) heatmap_loss = heatmap_loss * scale_factor if self._heatmap_weight is not None: # Create mask for scaling loss # Add 1.0 for saving values that are equal to 0 (approximately equal to 0) weight_mask = train_heatmap * self._heatmap_weight + self.__IDENTITY heatmap_loss = heatmap_loss * weight_mask heatmap_losses.append(tf.reduce_sum(heatmap_loss)) # The original repo takes mean over the sums of the losses self._paf_loss = tf.reduce_mean(paf_losses) self._heatmap_loss = tf.reduce_mean(heatmap_losses) loss = self._heatmap_loss * self._heatmap_scale + \ self._paf_loss * self._paf_scale # For Tensorboard super().track_loss(self._paf_loss, PETrainer.PAF_LOSS) super().track_loss(self._heatmap_loss, PETrainer.HEATMAP_LOSS) return loss