release code

Dassl.ProGrad.pytorch/dassl/modeling/backbone/cnn_digitsdg.py

@@ -0,0 +1,61 @@
+import torch.nn as nn
+from torch.nn import functional as F
+
+from dassl.utils import init_network_weights
+
+from .build import BACKBONE_REGISTRY
+from .backbone import Backbone
+
+
+class Convolution(nn.Module):
+
+    def __init__(self, c_in, c_out):
+        super().__init__()
+        self.conv = nn.Conv2d(c_in, c_out, 3, stride=1, padding=1)
+        self.relu = nn.ReLU(True)
+
+    def forward(self, x):
+        return self.relu(self.conv(x))
+
+
+class ConvNet(Backbone):
+
+    def __init__(self, c_hidden=64):
+        super().__init__()
+        self.conv1 = Convolution(3, c_hidden)
+        self.conv2 = Convolution(c_hidden, c_hidden)
+        self.conv3 = Convolution(c_hidden, c_hidden)
+        self.conv4 = Convolution(c_hidden, c_hidden)
+
+        self._out_features = 2**2 * c_hidden
+
+    def _check_input(self, x):
+        H, W = x.shape[2:]
+        assert (
+            H == 32 and W == 32
+        ), "Input to network must be 32x32, " "but got {}x{}".format(H, W)
+
+    def forward(self, x):
+        self._check_input(x)
+        x = self.conv1(x)
+        x = F.max_pool2d(x, 2)
+        x = self.conv2(x)
+        x = F.max_pool2d(x, 2)
+        x = self.conv3(x)
+        x = F.max_pool2d(x, 2)
+        x = self.conv4(x)
+        x = F.max_pool2d(x, 2)
+        return x.view(x.size(0), -1)
+
+
+@BACKBONE_REGISTRY.register()
+def cnn_digitsdg(**kwargs):
+    """
+    This architecture was used for DigitsDG dataset in:
+
+        - Zhou et al. Deep Domain-Adversarial Image Generation
+        for Domain Generalisation. AAAI 2020.
+    """
+    model = ConvNet(c_hidden=64)
+    init_network_weights(model, init_type="kaiming")
+    return model