Repository: pytorch
Issue: Domain Specific Batch Normalization
Labels: triaged,enhancement

## 🚀 Feature
<!-- A clear and concise description of the feature proposal -->
Domain-specific Batch Normalization ([paper](https://openaccess.thecvf.com/content_CVPR_2019/html/Chang_Domain-Specific_Batch_Normalization_for_Unsupervised_Domain_Adaptation_CVPR_2019_paper.html),  [github](https://github.com/wgchang/DSBN))

## Motivation
DSBN is a simple and powerful method in the domain adaptation method. ([paper](https://openaccess.thecvf.com/content_CVPR_2019/html/Chang_Domain-Specific_Batch_Normalization_for_Unsupervised_Domain_Adaptation_CVPR_2019_paper.html),  [github](https://github.com/wgchang/DSBN))

<!-- Please outline the motivation for the proposal. Is your feature request related to a problem? e.g., I'm always frustrated when [...]. If this is related to another GitHub issue, please link here too -->

## Pitch
<!-- A clear and concise description of what you want to happen. -->
It would be nice to add **a DSBN class and its methods** that enable data obtained from several different domains to be forwarded to the batch normalization layer corresponding to the specified domain.

I think it is necessary to implement the **set_domain_index** function that allows you to select which domain index batch norm layer to use in each network before forwarding data batches of each domain.

In addition, it would be nice to have a **convert_domain_specific_batchnorm** function that can utilize the pretrained BN parameters of the existing network by copying its parameters as many as the number of domains.

The code below is the code for how to use the DSBN described above. If there is a better way to implement it, it would be nice to give me your opinion.
<example usage>
```python
import torch.nn as nn
from torch.nn import DomainSpecificBatchNorm, DomainSpecificBatchNorm2d, set_domain_index

net = nn.Sequential(nn.Conv2d(3,10,3,bias=False), nn.BatchNorm2d(10), nn.ReLU(True),
                    nn.Conv2d(10,10,3,bias=False), nn.BatchNorm2d(10), nn.ReLU(True),
                    nn.Conv2d(10,2,3,bias=False))

dsbn_net = DomainSpecificBatchNorm.convert_domain_specific_batchnorm(net, num_domains=4)

d0_batch, d1_batch, d2_batch, d3_batch = # <... data loader of 4 domains ...>

set_domain_index(dsbn_net, 0)
d0_output = dsbn_net(d0_batch)

# <...>
set_domain_index(dsbn_net, 3)
d3_output = dsbn_net(d3_batch)
```

I would like to contribute by implementing DSBN in the pytorch repository.
The code below is a simple DSBN code by referring to the implementation of BatchNorm.
```python
class _DomainSpecificBatchNorm(nn.ModuleList):
    def __init__(self, num_features, num_domains=1, eps=1e-5, momentum=0.1, affine=True,
                 track_running_stats=True):
        super().__init__()
        self._domain_index = 0
        self._num_domains = num_domains
        self._set_batch_norm(num_features, num_domains, eps=1e-5, momentum=0.1, affine=True,
                             track_running_stats=True)

    def reset_running_stats(self):
        for m in self:
            if m.track_running_stats:
                m.running_mean.zero_()
                m.running_var.fill_(1)
                m.num_batches_tracked.zero_()

    def reset_parameters(self):
        for m in self:
            m.reset_running_stats()
            if m.affine:
                init.ones_(m.weight)
                init.zeros_(m.bias)

    def __repr__(self):
        return self._get_name() + '({num_features}, num_domains={num_domains}, eps={eps}, momentum={momentum}, affine={affine}, ' \
            'track_running_stats={track_running_stats})[{domain_index}]'.format(
            num_domains=self._num_domains, domain_index=self.domain_index, **self[0].__dict__)

    @property
    def num_domains(self):
        return self._num_domains

    @property
    def domain_index(self):
        return self._domain_index

    @domain_index.setter
    def domain_index(self, value):
        if value in range(self._num_domains):
            self._domain_index = value
        else:
            raise IndexError(f"Invalid domain index {value}")

    def forward(self, input):
        return self[self._domain_index](input)

    def _set_batch_norm(self, num_features, num_domains, eps=1e-5, momentum=0.1, affine=True,
                        track_running_stats=True):
        raise NotImplementedError

    @classmethod
    def set_domain_index(cls, modules, domain_index):
        if not isinstance(modules, (tuple, list)):
            modules = [modules, ]
        for module in modules:
            if isinstance(module, _DomainSpecificBatchNorm):
                module.domain_index = domain_index
            for name, child in module.named_children():
                cls.set_domain_index(child, domain_index)


class DomainSpecificBatchNorm1d(_DomainSpecificBatchNorm):
    def _set_batch_norm(self, num_features, num_domains, eps=1e-5, momentum=0.1, affine=True,
                        track_running_stats=True):
        self.extend([nn.BatchNorm1d(num_features, eps, momentum, affine,
                                    track_running_stats) for _ in range(self.num_domains)])


class DomainSpecificBatchNorm2d(_DomainSpecificBatchNorm):
    def _set_batch_norm(self, num_features, num_domains, eps=1e-5, momentum=0.1, affine=True,
                        track_running_stats=True):
        self.extend(
            [nn.BatchNorm2d(num_features, eps, momentum, affine, track_running_stats) for _ in range(self.num_domains)])


class DomainSpecificBatchNorm3d(_DomainSpecificBatchNorm):
    def _set_batch_norm(self, num_features, num_domains, eps=1e-5, momentum=0.1, affine=True,
                        track_running_stats=True):
        self.extend(
            [nn.BatchNorm3d(num_features, eps, momentum, affine, track_running_stats) for _ in range(self.num_domains)])


class SyncDomainSpecificBatchNorm(_DomainSpecificBatchNorm):
    def _set_batch_norm(self, num_features, num_domains, eps=1e-5, momentum=0.1, affine=True,
                        track_running_stats=True):
        self.extend(
            [nn.SyncBatchNorm(num_features, eps, momentum, affine, track_running_stats) for _ in range(self.num_domains)])


set_domain_index = _DomainSpecificBatchNorm.set_domain_index
```