diff --git a/ppcls/loss/__init__.py b/ppcls/loss/__init__.py index 90e9abf08..0e5070cc9 100644 --- a/ppcls/loss/__init__.py +++ b/ppcls/loss/__init__.py @@ -14,7 +14,7 @@ from .msmloss import MSMLoss from .npairsloss import NpairsLoss from .trihardloss import TriHardLoss from .triplet import TripletLoss, TripletLossV2 -from .tripletangularmarginloss import TripletAngularMarginLoss +from .tripletangularmarginloss import TripletAngularMarginLoss, TripletAngularMarginLoss_XBM from .supconloss import SupConLoss from .pairwisecosface import PairwiseCosface from .dmlloss import DMLLoss diff --git a/ppcls/loss/tripletangularmarginloss.py b/ppcls/loss/tripletangularmarginloss.py index 3a91d2d49..30b2d3138 100644 --- a/ppcls/loss/tripletangularmarginloss.py +++ b/ppcls/loss/tripletangularmarginloss.py @@ -18,6 +18,7 @@ from __future__ import print_function import paddle import paddle.nn as nn +from ppcls.loss.xbm import CrossBatchMemory class TripletAngularMarginLoss(nn.Layer): @@ -113,3 +114,128 @@ class TripletAngularMarginLoss(nn.Layer): ) * self.absolute_loss_weight + loss.mean() return {"TripletAngularMarginLoss": loss} + + +class TripletAngularMarginLoss_XBM(TripletAngularMarginLoss): + """TripletAngularMarginLoss combined with CrossBatchMemory + + Args: + start_iter: (int): from which step CrossBatchMemory is enabled + xbm_size: (int): Size of CrossBatchMemory + xbm_weight: (float): Weight of CrossBatchMemory loss + feat_dim: (int): Channels of features in CrossBatchMemory + margin (float, optional): angular margin. Defaults to 0.5. + normalize_feature (bool, optional): whether to apply L2-norm in feature before computing distance(cos-similarity). Defaults to True. + reduction (str, optional): reducing option within an batch . Defaults to "mean". + add_absolute (bool, optional): whether add absolute loss within d(a,p) or d(a,n). Defaults to False. + absolute_loss_weight (float, optional): weight for absolute loss. Defaults to 1.0. + ap_value (float, optional): weight for d(a, p). Defaults to 0.9. + an_value (float, optional): weight for d(a, n). Defaults to 0.5. + feature_from (str, optional): which key feature from. Defaults to "features". + """ + + def __init__(self, + start_iter: int, + xbm_size: int, + xbm_weight: float, + feat_dim: int, + margin=0.5, + normalize_feature=True, + reduction="mean", + add_absolute=False, + absolute_loss_weight=1.0, + ap_value=0.9, + an_value=0.5, + feature_from="features"): + super(TripletAngularMarginLoss_XBM, self).__init__( + margin, normalize_feature, reduction, add_absolute, + absolute_loss_weight, ap_value, an_value, feature_from) + self.start_iter = start_iter + self.xbm = CrossBatchMemory(xbm_size, feat_dim) + self.xbm_weight = xbm_weight + self.inf = 10 # 10 is big enough as inf for cos-similarity + self.register_buffer("iter", paddle.to_tensor(0, dtype="int64")) + + def forward(self, input, target): + """ + Args: + inputs: feature matrix with shape (batch_size, feat_dim) + target: ground truth labels with shape (num_classes) + """ + feats = input[self.feature_from] + if self.normalize_feature: + feats = nn.functional.normalize(feats, p=2, axis=1) + + labels = target + if labels.ndim >= 2 and labels.shape[-1] == 1: + labels = paddle.squeeze(labels, axis=[-1]) + + loss = self._compute_loss(feats, labels, feats, labels) + + # XBM loss below + self.iter += 1 + if self.iter.item() > self.start_iter: + self.xbm.enqueue_dequeue(feats.detach(), labels.detach()) + xbm_feats, xbm_labels = self.xbm.get() + xbm_loss = self._compute_loss(feats, labels, xbm_feats, xbm_labels) + loss = loss + self.xbm_weight * xbm_loss + + return {"TripletAngularMarginLoss_XBM": loss} + + def _masked_max(self, tensor, mask, axis): + masked = paddle.multiply(tensor, mask.astype(tensor.dtype)) + neg_inf = paddle.zeros_like(tensor) + neg_inf.stop_gradient = True + neg_inf[paddle.logical_not(mask)] = -self.inf + return paddle.max(masked + neg_inf, axis=axis, keepdim=True) + + def _masked_min(self, tensor, mask, axis): + masked = paddle.multiply(tensor, mask.astype(tensor.dtype)) + pos_inf = paddle.zeros_like(tensor) + pos_inf.stop_gradient = True + pos_inf[paddle.logical_not(mask)] = self.inf + return paddle.min(masked + pos_inf, axis=axis, keepdim=True) + + def _compute_loss(self, + inputs_q: paddle.Tensor, + targets_q: paddle.Tensor, + inputs_k: paddle.Tensor, + targets_k: paddle.Tensor) -> paddle.Tensor: + Q = inputs_q.shape[0] + K = inputs_k.shape[0] + + # compute distance(cos-similarity) + dist = paddle.matmul(inputs_q, inputs_k.t()) # [Q, K] + + # hard negative mining + is_pos = paddle.expand(paddle.unsqueeze(targets_q, 1), (Q, K)).equal( + paddle.expand(paddle.unsqueeze(targets_k, 1), + (K, Q)).t()) # [Q, K] + is_neg = paddle.expand(paddle.unsqueeze(targets_q, 1), + (Q, K)).not_equal( + paddle.expand( + paddle.unsqueeze(targets_k, 1), + (K, Q)).t()) # [Q, K] + + dist_ap = self._masked_min(dist, is_pos, axis=1) # [Q, ] + dist_an = self._masked_max(dist, is_neg, axis=1) # [Q, ] + + # Compute ranking hinge loss + y = paddle.ones_like(dist_an) + loss = self.ranking_loss(dist_ap, dist_an, y) + + if self.add_absolute: + absolut_loss_ap = self.ap_value - dist_ap + absolut_loss_ap = paddle.where(absolut_loss_ap > 0, + absolut_loss_ap, + paddle.zeros_like(absolut_loss_ap)) + + absolut_loss_an = dist_an - self.an_value + absolut_loss_an = paddle.where(absolut_loss_an > 0, + absolut_loss_an, + paddle.ones_like(absolut_loss_an)) + + loss = (absolut_loss_an.mean() + absolut_loss_ap.mean() + ) * self.absolute_loss_weight + loss.mean() + + return loss diff --git a/ppcls/loss/xbm.py b/ppcls/loss/xbm.py index 23656bf0a..b9f7fa34d 100644 --- a/ppcls/loss/xbm.py +++ b/ppcls/loss/xbm.py @@ -21,7 +21,7 @@ from typing import Tuple import paddle -class CrossBatchMemory(object): +class CrossBatchMemory(nn.Layer): """ CrossBatchMemory Implementation. refer to "Cross-Batch Memory for Embedding Learning". @@ -33,10 +33,18 @@ class CrossBatchMemory(object): """ def __init__(self, size: int, embedding_size: int): + super().__init__() self.size = size self.embedding_size = embedding_size - self.feats = paddle.zeros([self.size, self.embedding_size]) - self.targets = paddle.zeros([self.size, ], dtype="int64") + + # initialize and register feature queue for resume training + feats = paddle.zeros([self.size, self.embedding_size]) + self.register_buffer("feats", feats) + + # initialize and register label queue for resume training + targets = paddle.zeros([self.size, ], dtype="int64") + self.register_buffer("targets", targets) + self.ptr = 0 # self.accumulated_size = 0 @@ -74,3 +82,8 @@ class CrossBatchMemory(object): self.targets[self.ptr:self.ptr + input_size] = targets self.ptr += input_size # self.accumulated_size += input_size + + def forward(self, *kargs, **kwargs): + raise NotImplementedError( + "CrossBatchMemory module is for memory-bank, forward method is not needed" + )