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)