add reweight / normalization package

Author: JunguangJiang

dglib/generalization/__init__.py

@@ -5,7 +5,7 @@
 from typing import Optional, Tuple
 import torch
 import torch.nn as nn
-from common.modules.classifier import Classifier as ClassifierBase
+from tllib.modules.classifier import Classifier as ClassifierBase
 
 
 class ImageClassifier(ClassifierBase):

dglib/generalization/irm.py

@@ -22,7 +22,7 @@
 from tllib.modules.classifier import Classifier
 from tllib.modules.entropy import entropy
 from tllib.modules.domain_discriminator import DomainDiscriminator
-from tllib.alignment.iwan import ImportanceWeightModule, ImageClassifier
+from tllib.reweight.iwan import ImportanceWeightModule, ImageClassifier
 from tllib.alignment.dann import DomainAdversarialLoss
 from tllib.utils.data import ForeverDataIterator
 from tllib.utils.metric import accuracy

dglib/modules/classifier.py

@@ -22,7 +22,7 @@
 from tllib.modules.domain_discriminator import DomainDiscriminator
 from tllib.modules.classifier import Classifier
 from tllib.alignment.dann import DomainAdversarialLoss, ImageClassifier
-from tllib.alignment.pada import AutomaticUpdateClassWeightModule
+from tllib.reweight.pada import AutomaticUpdateClassWeightModule
 from tllib.utils.data import ForeverDataIterator
 from tllib.utils.metric import accuracy
 from tllib.utils.meter import AverageMeter, ProgressMeter

examples/domain_adaptation/partial_domain_adaptation/iwan.py

@@ -1,4 +1,5 @@
 """
+Adapted from https://github.com/facebookresearch/DomainBed
 @author: Baixu Chen
 @contact: [email protected]
 """
@@ -19,9 +20,9 @@
 import torch.nn.functional as F
 
 sys.path.append('../../..')
-from dglib.generalization.groupdro import AutomaticUpdateDomainWeightModule
 from dglib.modules.sampler import RandomDomainSampler
 from dglib.modules.classifier import ImageClassifier as Classifier
+from tllib.reweight.groupdro import AutomaticUpdateDomainWeightModule
 from tllib.utils.data import ForeverDataIterator
 from tllib.utils.metric import accuracy
 from tllib.utils.meter import AverageMeter, ProgressMeter

examples/domain_adaptation/partial_domain_adaptation/pada.py

@@ -1,4 +1,5 @@
 """
+Adapted from https://github.com/facebookresearch/DomainBed
 @author: Baixu Chen
 @contact: [email protected]
 """
@@ -17,9 +18,9 @@
 from torch.optim.lr_scheduler import CosineAnnealingLR
 from torch.utils.data import DataLoader
 import torch.nn.functional as F
+import torch.autograd as autograd
 
 sys.path.append('../../..')
-from dglib.generalization.irm import InvariancePenaltyLoss
 from dglib.modules.sampler import RandomDomainSampler
 from dglib.modules.classifier import ImageClassifier as Classifier
 from tllib.utils.data import ForeverDataIterator
@@ -34,6 +35,35 @@
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 
+class InvariancePenaltyLoss(nn.Module):
+    r"""Invariance Penalty Loss from `Invariant Risk Minimization <https://arxiv.org/pdf/1907.02893.pdf>`_.
+    We adopt implementation from `DomainBed <https://github.com/facebookresearch/DomainBed>`_. Given classifier
+    output :math:`y` and ground truth :math:`labels`, we split :math:`y` into two parts :math:`y_1, y_2`, corresponding
+    labels are :math:`labels_1, labels_2`. Next we calculate cross entropy loss with respect to a dummy classifier
+    :math:`w`, resulting in :math:`grad_1, grad_2` . Invariance penalty is then :math:`grad_1*grad_2`.
+
+    Inputs:
+        - y: predictions from model
+        - labels: ground truth
+
+    Shape:
+        - y: :math:`(N, C)` where C means the number of classes.
+        - labels: :math:`(N, )` where N mean mini-batch size
+    """
+
+    def __init__(self):
+        super(InvariancePenaltyLoss, self).__init__()
+        self.scale = torch.tensor(1.).requires_grad_()
+
+    def forward(self, y: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
+        loss_1 = F.cross_entropy(y[::2] * self.scale, labels[::2])
+        loss_2 = F.cross_entropy(y[1::2] * self.scale, labels[1::2])
+        grad_1 = autograd.grad(loss_1, [self.scale], create_graph=True)[0]
+        grad_2 = autograd.grad(loss_2, [self.scale], create_graph=True)[0]
+        penalty = torch.sum(grad_1 * grad_2)
+        return penalty
+
+
 def main(args: argparse.Namespace):
     logger = CompleteLogger(args.log, args.phase)
     print(args)

examples/domain_generalization/image_classification/groupdro.py

@@ -19,7 +19,7 @@
 import torch.nn.functional as F
 
 sys.path.append('../../..')
-import dglib.generalization.mixstyle.models as models
+import tllib.modules.normalization.mixstyle.models as models
 from dglib.modules.sampler import RandomDomainSampler
 from dglib.modules.classifier import ImageClassifier as Classifier
 from tllib.utils.data import ForeverDataIterator

examples/domain_generalization/image_classification/irm.py

@@ -18,8 +18,8 @@
 from torch.utils.data import DataLoader
 
 sys.path.append('../../..')
-from dglib.generalization.mixstyle.sampler import RandomDomainMultiInstanceSampler
-import dglib.generalization.mixstyle.models as models
+from tllib.modules.normalization.mixstyle.sampler import RandomDomainMultiInstanceSampler
+import tllib.modules.normalization.mixstyle.models as models
 from tllib.vision.models.reid.identifier import ReIdentifier
 from tllib.vision.models.reid.loss import CrossEntropyLossWithLabelSmooth, SoftTripletLoss
 import tllib.vision.datasets.reid as datasets

examples/domain_generalization/image_classification/mixstyle.py

@@ -17,7 +17,7 @@
 import torch.nn.functional as F
 
 sys.path.append('../../..')
-from tllib.regularization.stochnorm import convert_model
+from tllib.modules.normalization.stochnorm import convert_model
 from tllib.modules.classifier import Classifier
 from tllib.utils.metric import accuracy
 from tllib.utils.meter import AverageMeter, ProgressMeter

examples/domain_generalization/re_identification/mixstyle.py

@@ -1,6 +1,7 @@
 # Copy from https://github.com/CoinCheung/pytorch-loss
 import torch.nn as nn
 import torch
+import torch.nn.functional as F
 
 
 # version 1: use torch.autograd
@@ -45,3 +46,32 @@ def forward(self, input, target):
             loss = loss.sum()
 
         return loss
+
+
+class KnowledgeDistillationLoss(nn.Module):
+    """Knowledge Distillation Loss.
+
+    Args:
+        T (double): Temperature. Default: 1.
+        reduction (str, optional): Specifies the reduction to apply to the output:
+          ``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``: no reduction will be applied,
+          ``'mean'``: the sum of the output will be divided by the number of
+          elements in the output, ``'sum'``: the output will be summed. Default: ``'batchmean'``
+
+    Inputs:
+        - y_student (tensor): logits output of the student
+        - y_teacher (tensor): logits output of the teacher
+
+    Shape:
+        - y_student: (minibatch, `num_classes`)
+        - y_teacher: (minibatch, `num_classes`)
+
+    """
+    def __init__(self, T=1., reduction='batchmean'):
+        super(KnowledgeDistillationLoss, self).__init__()
+        self.T = T
+        self.kl = nn.KLDivLoss(reduction=reduction)
+
+    def forward(self, y_student, y_teacher):
+        """"""
+        return self.kl(F.log_softmax(y_student / self.T, dim=-1), F.softmax(y_teacher / self.T, dim=-1))