text: import torch
  import torch.nn as nn
  from torch.nn import functional as F
  
  from dassl.optim import build_optimizer, build_lr_scheduler
  from dassl.utils import count_num_param
  from dassl.engine import TRAINER_REGISTRY, TrainerXU
  from dassl.engine.trainer import SimpleNet
  
  
  class PairClassifiers(nn.Module):
  
      def __init__(self, fdim, num_classes):
          super().__init__()
          self.c1 = nn.Linear(fdim, num_classes)
          self.c2 = nn.Linear(fdim, num_classes)
  
      def forward(self, x):
          z1 = self.c1(x)
          if not self.training:
              return z1
          z2 = self.c2(x)
          return z1, z2
  
  
  @TRAINER_REGISTRY.register()
  class M3SDA(TrainerXU):
      """Moment Matching for Multi-Source Domain Adaptation.
  
      https://arxiv.org/abs/1812.01754.
      """
  
      def __init__(self, cfg):
          super().__init__(cfg)
          n_domain = cfg.DATALOADER.TRAIN_X.N_DOMAIN
          batch_size = cfg.DATALOADER.TRAIN_X.BATCH_SIZE
          if n_domain <= 0:
              n_domain = self.dm.num_source_domains
          self.split_batch = batch_size // n_domain
          self.n_domain = n_domain
  
          self.n_step_F = cfg.TRAINER.M3SDA.N_STEP_F
          self.lmda = cfg.TRAINER.M3SDA.LMDA
  
      def check_cfg(self, cfg):
          assert cfg.DATALOADER.TRAIN_X.SAMPLER == 'RandomDomainSampler'
          assert not cfg.DATALOADER.TRAIN_U.SAME_AS_X
  
      def build_model(self):
          cfg = self.cfg
  
          print('Building F')
          self.F = SimpleNet(cfg, cfg.MODEL, 0)
          self.F.to(self.device)
          print('# params: {:,}'.format(count_num_param(self.F)))
          self.optim_F = build_optimizer(self.F, cfg.OPTIM)
          self.sched_F = build_lr_scheduler(self.optim_F, cfg.OPTIM)
          self.register_model('F', self.F, self.optim_F, self.sched_F)
          fdim = self.F.fdim
  
          print('Building C')
          self.C = nn.ModuleList(
              [
                  PairClassifiers(fdim, self.num_classes)
                  for _ in range(self.dm.num_source_domains)
              ]
          )
          self.C.to(self.device)
          print('# params: {:,}'.format(count_num_param(self.C)))
          self.optim_C = build_optimizer(self.C, cfg.OPTIM)
          self.sched_C = build_lr_scheduler(self.optim_C, cfg.OPTIM)
          self.register_model('C', self.C, self.optim_C, self.sched_C)
  
      def forward_backward(self, batch_x, batch_u):
          parsed = self.parse_batch_train(batch_x, batch_u)
          input_x, label_x, domain_x, input_u = parsed
  
          input_x = torch.split(input_x, self.split_batch, 0)
          label_x = torch.split(label_x, self.split_batch, 0)
          domain_x = torch.split(domain_x, self.split_batch, 0)
          domain_x = [d[0].item() for d in domain_x]
  
          # Step A
          loss_x = 0
          feat_x = []
  
          for x, y, d in zip(input_x, label_x, domain_x):
              f = self.F(x)
              z1, z2 = self.C[d](f)
              loss_x += F.cross_entropy(z1, y) + F.cross_entropy(z2, y)
  
              feat_x.append(f)
  
          loss_x /= self.n_domain
  
          feat_u = self.F(input_u)
          loss_msda = self.moment_distance(feat_x, feat_u)
  
          loss_step_A = loss_x + loss_msda * self.lmda
          self.model_backward_and_update(loss_step_A)
  
          # Step B
          with torch.no_grad():
              feat_u = self.F(input_u)
  
          loss_x, loss_dis = 0, 0
  
          for x, y, d in zip(input_x, label_x, domain_x):
              with torch.no_grad():
                  f = self.F(x)
              z1, z2 = self.C[d](f)
              loss_x += F.cross_entropy(z1, y) + F.cross_entropy(z2, y)
  
              z1, z2 = self.C[d](feat_u)
              p1 = F.softmax(z1, 1)
              p2 = F.softmax(z2, 1)
              loss_dis += self.discrepancy(p1, p2)
  
          loss_x /= self.n_domain
          loss_dis /= self.n_domain
  
          loss_step_B = loss_x - loss_dis
          self.model_backward_and_update(loss_step_B, 'C')
  
          # Step C
          for _ in range(self.n_step_F):
              feat_u = self.F(input_u)
  
              loss_dis = 0
  
              for d in domain_x:
                  z1, z2 = self.C[d](feat_u)
                  p1 = F.softmax(z1, 1)
                  p2 = F.softmax(z2, 1)
                  loss_dis += self.discrepancy(p1, p2)
  
              loss_dis /= self.n_domain
              loss_step_C = loss_dis
  
              self.model_backward_and_update(loss_step_C, 'F')
  
          output_dict = {
              'loss_step_A': loss_step_A.item(),
              'loss_step_B': loss_step_B.item(),
              'loss_step_C': loss_step_C.item(),
              'lr': self.optim_F.param_groups[0]['lr']
          }
  
          if (self.batch_idx + 1) == self.num_batches:
              self.update_lr()
  
          return output_dict
  
      def moment_distance(self, x, u):
          # x (list): a list of feature matrix.
          # u (torch.Tensor): feature matrix.
          x_mean = [xi.mean(0) for xi in x]
          u_mean = u.mean(0)
          dist1 = self.pairwise_distance(x_mean, u_mean)
  
          x_var = [xi.var(0) for xi in x]
          u_var = u.var(0)
          dist2 = self.pairwise_distance(x_var, u_var)
  
          return (dist1+dist2) / 2
  
      def pairwise_distance(self, x, u):
          # x (list): a list of feature vector.
          # u (torch.Tensor): feature vector.
          dist = 0
          count = 0
  
          for xi in x:
              dist += self.euclidean(xi, u)
              count += 1
  
          for i in range(len(x) - 1):
              for j in range(i + 1, len(x)):
                  dist += self.euclidean(x[i], x[j])
                  count += 1
  
          return dist / count
  
      def euclidean(self, input1, input2):
          return ((input1 - input2)**2).sum().sqrt()
  
      def discrepancy(self, y1, y2):
          return (y1 - y2).abs().mean()
  
      def parse_batch_train(self, batch_x, batch_u):
          input_x = batch_x['img']
          label_x = batch_x['label']
          domain_x = batch_x['domain']
          input_u = batch_u['img']
  
          input_x = input_x.to(self.device)
          label_x = label_x.to(self.device)
          input_u = input_u.to(self.device)
  
          return input_x, label_x, domain_x, input_u
  
      def model_inference(self, input):
          f = self.F(input)
          p = 0
          for C_i in self.C:
              z = C_i(f)
              p += F.softmax(z, 1)
          p = p / len(self.C)
          return p
  