|
|
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 |
|
|
|
