def __init__(self,
in_channels,
out_channels,
hidden_channels,
downsample=False,
spectral_norm=True):
super().__init__()
self.in_channels = in_channels
self.out_channels = out_channels
self.hidden_channels = hidden_channels \
if hidden_channels is not None else in_channels
self.downsample = downsample
self.learnable_sc = (in_channels != out_channels) or downsample
self.spectral_norm = spectral_norm
if self.spectral_norm:
self.c1 = SNConv2d(self.in_channels, self.hidden_channels, 3, 1, 1)
self.c2 = SNConv2d(self.hidden_channels, self.out_channels, 3, 1,
1)
else:
self.c1 = nn.Conv2d(self.in_channels, self.hidden_channels, 3, 1,
1)
self.c2 = nn.Conv2d(self.hidden_channels, self.out_channels, 3, 1,
1)
self.average_pooling = nn.AvgPool2d(2)
self.activation = nn.ReLU(inplace=True)
# Shortcut layer
if self.learnable_sc:
if self.spectral_norm:
self.c_sc = SNConv2d(in_channels, out_channels, 1, 1, 0)
else:
self.c_sc = nn.Conv2d(in_channels, out_channels, 1, 1, 0)
def __init__(self,
ndf=128,
loss_type="ns",
is_amp=False,
cutmix=False,
consistency=False,
warmup=50000,
**kwargs):
super().__init__(ndf=ndf, loss_type=loss_type, **kwargs)
self.is_amp = is_amp
self.cutmix = cutmix
self.consistency = consistency
# build encoder
self.down1 = DBlockOptimized(3, self.ndf)
self.down2 = DBlock(self.ndf, self.ndf, downsample=True)
self.down3 = DBlock(self.ndf, self.ndf, downsample=False)
self.down4 = DBlock(self.ndf, self.ndf, downsample=False)
self.down_act = nn.ReLU(inplace=True)
self.down_l5 = SNLinear(self.ndf, 1)
# self-attention
self.non_local_block = SelfAttention(self.ndf, spectral_norm=True)
# initialize weights
nn.init.xavier_uniform_(self.down_l5.weight.data, 1.0)
# build decoder
self.up1 = DBlockDecoder(self.ndf, self.ndf, upsample=False)
self.up2 = DBlockDecoder(self.ndf * 2, self.ndf, upsample=False)
self.up3 = DBlockDecoder(self.ndf * 2, self.ndf, upsample=True)
self.up4 = DBlockDecoder(self.ndf * 2, self.ndf, upsample=True)
self.up_act = nn.ReLU(inplace=True)
self.up_c5 = SNConv2d(self.ndf, 1, kernel_size=3, stride=1, padding=1)
# initialize weights
nn.init.xavier_uniform_(self.up_c5.weight.data, 1.0)
def __init__(self,
in_channels,
out_channels,
hidden_channels=None,
upsample=False,
concat_vector_dim=None,
spectral_norm=False):
super().__init__()
self.in_channels = in_channels
self.out_channels = out_channels
self.hidden_channels = hidden_channels if hidden_channels is not None else out_channels
self.learnable_sc = in_channels != out_channels or upsample
self.upsample = upsample
self.concat_vector_dim = concat_vector_dim
self.spectral_norm = spectral_norm
if self.spectral_norm:
self.c1 = SNConv2d(self.in_channels, self.hidden_channels, 3, 1, 1)
self.c2 = SNConv2d(self.hidden_channels, self.out_channels, 3, 1,
1)
else:
self.c1 = nn.Conv2d(self.in_channels, self.hidden_channels, 3, 1,
1)
self.c2 = nn.Conv2d(self.hidden_channels, self.out_channels, 3, 1,
1)
if self.num_classes == 0:
self.b1 = nn.BatchNorm2d(self.in_channels)
self.b2 = nn.BatchNorm2d(self.hidden_channels)
else:
self.b1 = ConditionalBatchNorm2d_with_skip_and_shared(
self.in_channels, self.concat_vector_dim)
self.b2 = ConditionalBatchNorm2d_with_skip_and_shared(
self.hidden_channels, self.concat_vector_dim)
self.activation = nn.ReLU(True)
# Shortcut layer
if self.learnable_sc:
if self.spectral_norm:
self.c_sc = SNConv2d(in_channels, out_channels, 1, 1, 0)
else:
self.c_sc = nn.Conv2d(in_channels, out_channels, 1, 1, 0)
def __init__(self, in_channels, out_channels, spectral_norm=True):
super().__init__()
self.in_channels = in_channels
self.out_channels = out_channels
self.spectral_norm = spectral_norm
# Build the layers
if self.spectral_norm:
self.c1 = SNConv2d(self.in_channels, self.out_channels, 3, 1, 1)
self.c2 = SNConv2d(self.out_channels, self.out_channels, 3, 1, 1)
self.c_sc = SNConv2d(self.in_channels, self.out_channels, 1, 1, 0)
else:
self.c1 = nn.Conv2d(self.in_channels, self.out_channels, 3, 1, 1)
self.c2 = nn.Conv2d(self.out_channels, self.out_channels, 3, 1, 1)
self.c_sc = nn.Conv2d(self.in_channels, self.out_channels, 1, 1, 0)
self.activation = nn.ReLU(True)
nn.init.xavier_uniform_(self.c1.weight.data, math.sqrt(2.0))
nn.init.xavier_uniform_(self.c2.weight.data, math.sqrt(2.0))
nn.init.xavier_uniform_(self.c_sc.weight.data, 1.0)
def __init__(self,
in_channels,
out_channels,
hidden_channels=None,
downsample=False,
spectral_norm=True):
super().__init__()
self.in_channels = in_channels
self.out_channels = out_channels
self.hidden_channels = hidden_channels if hidden_channels is not None else in_channels
self.downsample = downsample
self.learnable_sc = (in_channels != out_channels) or downsample
self.spectral_norm = spectral_norm
# Build the layers
if self.spectral_norm:
self.c1 = SNConv2d(self.in_channels, self.hidden_channels, 3, 1, 1)
self.c2 = SNConv2d(self.hidden_channels, self.out_channels, 3, 1,
1)
else:
self.c1 = nn.Conv2d(self.in_channels, self.hidden_channels, 3, 1,
1)
self.c2 = nn.Conv2d(self.hidden_channels, self.out_channels, 3, 1,
1)
self.activation = nn.ReLU(True)
nn.init.xavier_uniform_(self.c1.weight.data, math.sqrt(2.0))
nn.init.xavier_uniform_(self.c2.weight.data, math.sqrt(2.0))
# Shortcut layer
if self.learnable_sc:
if self.spectral_norm:
self.c_sc = SNConv2d(in_channels, out_channels, 1, 1, 0)
else:
self.c_sc = nn.Conv2d(in_channels, out_channels, 1, 1, 0)
nn.init.xavier_uniform_(self.c_sc.weight.data, 1.0)
def __init__(self,
in_channels,
out_channels,
hidden_channels=None,
upsample=False,
num_classes=0,
spectral_norm=False,
nz=128):
super().__init__()
self.in_channels = in_channels
self.out_channels = out_channels
self.hidden_channels = hidden_channels if hidden_channels is not None else out_channels
self.learnable_sc = in_channels != out_channels or upsample
self.upsample = upsample
self.num_classes = num_classes
self.spectral_norm = spectral_norm
self.nz = nz
# build the layers
if self.spectral_norm:
self.c1 = SNConv2d(self.in_channels,
self.hidden_channels,
kernel_size=3,
stride=1,
padding=1)
self.c2 = SNConv2d(self.in_channels,
self.hidden_channels,
kernel_size=3,
stride=1,
padding=1)
else:
self.c1 = nn.Conv2d(self.in_channels,
self.hidden_channels,
kernel_size=3,
stride=1,
padding=1)
self.c2 = nn.Conv2d(self.hidden_channels,
self.out_channels,
kernel_size=3,
stride=1,
padding=1)
if self.num_classes == 0:
self.b1 = SBN(self.in_channels, self.nz)
self.b2 = SBN(self.hidden_channels, self.nz)
else:
self.b1 = SCBN(self.in_channels, self.num_classes, self.nz)
self.b2 = SCBN(self.hidden_channels, self.num_classes, self.nz)
self.activation = nn.ReLU(inplace=True)
nn.init.xavier_uniform_(self.c1.weight.data, math.sqrt(2.0))
nn.init.xavier_uniform_(self.c2.weight.data, math.sqrt(2.0))
if self.learnable_sc:
if self.spectral_norm:
self.c_sc = SNConv2d(self.in_channels,
self.out_channels,
kernel_size=1,
stride=1,
padding=0)
else:
self.c_sc = nn.Conv2d(self.in_channels,
self.out_channels,
kernel_size=1,
stride=1,
padding=0)
nn.init.xavier_uniform_(self.c_sc.weight.data, 1.0)
def __init__(self,
in_channels,
out_channels,
hidden_channels=None,
upsample=False,
num_classes=0,
spectral_norm=False):
super().__init__()
self.in_channels = in_channels
self.out_channels = out_channels
self.hidden_channels = hidden_channels if hidden_channels is not None else out_channels
self.learnable_sc = in_channels != out_channels or upsample
self.upsample = upsample
self.num_classes = num_classes
self.spectral_norm = spectral_norm
# Build the layers
if self.spectral_norm:
self.c1 = SNConv2d(self.in_channels,
self.hidden_channels,
3,
1,
padding=1)
self.c2 = SNConv2d(self.hidden_channels,
self.out_channels,
3,
1,
padding=1)
else:
self.c1 = nn.Conv2d(self.in_channels,
self.hidden_channels,
3,
1,
padding=1)
self.c2 = nn.Conv2d(self.hidden_channels,
self.out_channels,
3,
1,
padding=1)
if self.num_classes == 0:
self.b1 = nn.BatchNorm2d(self.in_channels)
self.b2 = nn.BatchNorm2d(self.hidden_channels)
else:
self.b1 = ConditionalBatchNorm2d(self.in_channels,
self.num_classes)
self.b2 = ConditionalBatchNorm2d(self.hidden_channels,
self.num_classes)
self.activation = nn.ReLU(True)
nn.init.xavier_uniform_(self.c1.weight.data, math.sqrt(2.0))
nn.init.xavier_uniform_(self.c2.weight.data, math.sqrt(2.0))
# Shortcut layer
if self.learnable_sc:
if self.spectral_norm:
self.c_sc = SNConv2d(in_channels,
out_channels,
1,
1,
padding=0)
else:
self.c_sc = nn.Conv2d(in_channels,
out_channels,
1,
1,
padding=0)
nn.init.xavier_uniform_(self.c_sc.weight.data, 1.0)