def make_layers(cfg, batch_norm=False,deconv=None, norm_type='none'):
layers = []
in_channels = 3
if not deconv:
for v in cfg:
if v == 'M':
layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
else:
conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
if batch_norm:
layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
else:
layers += [conv2d, nn.ReLU(inplace=True)]
in_channels = v
else:
for i,v in enumerate(cfg):
if v == 'M':
layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
else:
if in_channels==3:
conv2d = deconv(in_channels, v, kernel_size=3, padding=1, freeze=True, n_iter=15, sampling_stride=3)
else:
conv2d = deconv(in_channels, v, kernel_size=3, padding=1)
if batch_norm:
layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
else:
layers += [conv2d, nn.ReLU(inplace=True)]
in_channels = v
return nn.Sequential(*layers)
def conv1x1(in_planes, out_planes, stride=1, deconv=None):
"""1x1 convolution"""
if deconv:
return deconv(in_planes, out_planes, kernel_size=1, stride=stride)
else:
return nn.Conv2d(in_planes,
out_planes,
kernel_size=1,
stride=stride,
bias=False)
def __init__(self, conv_dim=64, noise_dim=32, init_zero_weights=False):
super(CycleGenerator, self).__init__()
self.conv_dim = conv_dim
self.noise_dim = noise_dim
# 1. Define the encoder part of the generator (that extracts features from the input image)
self.conv1 = conv(3, conv_dim, 4, padding=1, stride=2)
self.conv2 = conv(conv_dim, conv_dim * 2, 4, padding=1, stride=2)
# 2. Define the transformation part of the generator
self.resnet_block = ResnetBlock(conv_dim * 2)
# 3. Define the decoder part of the generator (that builds up the output image from features)
self.deconv1 = deconv(conv_dim * 2 + noise_dim,
conv_dim,
4,
padding=1,
stride=2)
self.deconv2 = deconv(conv_dim,
3,
4,
padding=1,
stride=2,
batch_norm=False)
def conv3x3(in_planes,
out_planes,
stride=1,
groups=1,
dilation=1,
deconv=None):
"""3x3 convolution with padding"""
if deconv:
return deconv(in_planes,
out_planes,
kernel_size=3,
stride=stride,
padding=dilation,
dilation=dilation) #
else:
return nn.Conv2d(in_planes,
out_planes,
kernel_size=3,
stride=stride,
padding=dilation,
bias=False,
dilation=dilation,
groups=groups) #
def __init__(self,
block,
layers,
num_classes=1000,
zero_init_residual=False,
groups=1,
width_per_group=64,
replace_stride_with_dilation=None,
norm_layer=None,
deconv=None,
channel_deconv=None):
super(ResNet, self).__init__()
if norm_layer is None:
norm_layer = nn.BatchNorm2d
self._norm_layer = norm_layer
self.inplanes = 64
self.dilation = 1
if replace_stride_with_dilation is None:
# each element in the tuple indicates if we should replace
# the 2x2 stride with a dilated convolution instead
replace_stride_with_dilation = [False, False, False]
if len(replace_stride_with_dilation) != 3:
raise ValueError("replace_stride_with_dilation should be None "
"or a 3-element tuple, got {}".format(
replace_stride_with_dilation))
self.groups = groups
self.base_width = width_per_group
if not deconv:
self.conv1 = nn.Conv2d(3,
self.inplanes,
kernel_size=7,
stride=2,
padding=3,
bias=False)
self.bn1 = norm_layer(self.inplanes)
else:
self.conv1 = deconv(3,
self.inplanes,
kernel_size=7,
stride=2,
padding=3)
# this line is really recent, take extreme care if the result is not good. Batch size really matters here.
if channel_deconv:
self.deconv1 = channel_deconv()
self.relu = nn.ReLU(inplace=True)
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
self.layer1 = self._make_layer(block, 64, layers[0], deconv=deconv)
self.layer2 = self._make_layer(block,
128,
layers[1],
stride=2,
dilate=replace_stride_with_dilation[0],
deconv=deconv)
self.layer3 = self._make_layer(block,
256,
layers[2],
stride=2,
dilate=replace_stride_with_dilation[1],
deconv=deconv)
self.layer4 = self._make_layer(block,
512,
layers[3],
stride=2,
dilate=replace_stride_with_dilation[2],
deconv=deconv)
self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
self.fc = nn.Linear(512 * block.expansion, num_classes)
for m in self.modules():
if isinstance(m, nn.Conv2d) or isinstance(m, DeConv2d):
nn.init.kaiming_normal_(m.weight,
mode='fan_out',
nonlinearity='relu')
elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
nn.init.constant_(m.weight, 1)
nn.init.constant_(m.bias, 0)
# Zero-initialize the last BN in each residual branch,
# so that the residual branch starts with zeros, and each residual block behaves like an identity.
# This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
if zero_init_residual:
for m in self.modules():
if isinstance(m, Bottleneck):
nn.init.constant_(m.bn3.weight, 0)
elif isinstance(m, BasicBlock):
nn.init.constant_(m.bn2.weight, 0)
