def __init__(self, arch: List[int], n_skip_max: int = 2, in_noise: int = 256, out_ch: int = 3, batchnorm_in_output: bool = False) -> None: super().__init__() self.n_skip_max = n_skip_max self.make_noise = xavier(nn.Linear(in_noise, 4 * 4 * arch[0])) in_ch = arch[0] blocks = [] lasts: List[int] = [] for i, out in enumerate(arch[1:]): mode = 'nearest' if i % 2 == 0 else 'bilinear' blocks.append(tnn.AutoGANGenBlock(in_ch, out, lasts, mode=mode)) lasts = ([out] + lasts)[:n_skip_max] in_ch = out self.blocks = nn.ModuleList(blocks) if batchnorm_in_output: self.to_rgb = nn.Sequential(nn.BatchNorm2d(arch[-1]), nn.ReLU(True), xavier(tnn.Conv3x3(arch[-1], out_ch))) else: self.to_rgb = nn.Sequential(nn.ReLU(True), xavier(tnn.Conv3x3(arch[-1], out_ch)))
def __init__(self, in_ch, out_ch, ks, stride=1, mul_factor=6): super(MBConv, self).__init__() self.in_ch = in_ch self.out_ch = out_ch self.ks = ks self.stride = stride self.factor = mul_factor hid = in_ch * mul_factor self.branch = tnn.CondSeq( tu.xavier(tnn.Conv1x1(in_ch, hid, bias=False)), nn.BatchNorm2d(hid), tnn.HardSwish(), tu.xavier( nn.Conv2d(hid, hid, ks, stride=stride, padding=ks // 2, groups=hid, bias=False)), nn.BatchNorm2d(hid), tnn.HardSwish(), tnn.SEBlock(hid, reduction=4), tu.xavier(tnn.Conv1x1(hid, out_ch)), nn.BatchNorm2d(out_ch)) self.shortcut = tnn.CondSeq() if stride != 1: self.shortcut.add_module( 'pool', nn.AvgPool2d(stride, stride, ceil_mode=True)) if in_ch != out_ch: self.shortcut.add_module('conv', tnn.Conv1x1(in_ch, out_ch, bias=False)) self.shortcut.add_module('bn', nn.BatchNorm2d(out_ch))
def __init__(self, in_ch: int, out_ch: int, noise_size: int, upsample: bool = False, n_layers: int = 2, equal_lr: bool = True): super().__init__() self.upsample_mode = 'bilinear' self.equal_lr = equal_lr dyn = equal_lr self.upsample = upsample inside = ModuleGraph(outputs=f'in_{n_layers}') for i in range(n_layers): conv = ModulatedConv(in_ch, noise_size, out_ch, kernel_size=3, padding=1, bias=True) kaiming(conv, dynamic=dyn, a=0.2) inside.add_operation(inputs=[f'in_{i}', 'w'], outputs=[f'conv_{i}'], name=f'conv_{i}', operation=conv) noise = Noise(out_ch, inplace=True, bias=False) inside.add_operation(inputs=[f'conv_{i}', f'noise_{i}'], operation=noise, name=f'plus_noise_{i}', outputs=[f'plus_noise_{i}']) inside.add_operation(inputs=[f'plus_noise_{i}'], operation=nn.LeakyReLU(0.2, True), outputs=[f'in_{i+1}'], name=f'relu_{i+1}') in_ch = out_ch if dyn: for m in inside: if isinstance(m, ModulatedConv): xavier(m.make_s, dynamic=True) self.inside = inside self.to_rgb = xavier(ModulatedConv(out_ch, noise_size, 3, kernel_size=1, padding=0, bias=True, demodulate=False), dynamic=dyn, a=0.2)
def __init__(self, channels, cond_channels, hidden, size=None, momentum=0.8): super(Spade2d, self).__init__(channels, momentum) self.initial = kaiming(Conv3x3(cond_channels, hidden, stride=2)) self.make_weight = xavier(Conv3x3(hidden, channels)) self.make_bias = xavier(Conv3x3(hidden, channels)) self.size = size
def to_spectral_norm(self) -> 'ResidualDiscrBlock': for m in self.modules(): if isinstance(m, nn.Conv2d): try: remove_weight_scale(m) except ValueError: pass nn.utils.spectral_norm(m) assert isinstance(self.branch.conv2, nn.Conv2d) xavier(self.branch.conv2, a=0.2) return self
def __init__(self, in_ch, reduction=16): super(SEBlock, self).__init__() self.proj = nn.Sequential( collections.OrderedDict([ ('pool', nn.AdaptiveAvgPool2d(1)), ('squeeze', kaiming(nn.Conv2d(in_ch, in_ch // reduction, 1))), ('relu', nn.ReLU(True)), ('excite', xavier(nn.Conv2d(in_ch // reduction, in_ch, 1))), ('attn', HardSigmoid()) ]))
def __init__(self, in_channels: int, noise_channels: int, *args, demodulate: bool = True, **kwargs): super(ModulatedConv, self).__init__(in_channels, *args, **kwargs) self.make_s = tu.xavier(nn.Linear(noise_channels, in_channels)) self.make_s.bias.data.fill_(1) self.demodulate = demodulate
def __init__(self, feat_extractor, feature_size, num_classes): super(Classifier, self).__init__() self.bone = feat_extractor self.head = nn.Sequential( nn.AdaptiveMaxPool2d(1), tnn.Reshape(feature_size), kaiming(nn.Linear(feature_size, feature_size)), nn.ReLU(inplace=True), xavier(nn.Linear(feature_size, num_classes)), )
def __init__(self, in_ch, num_classes, B=0): def ch(ch): return int(ch * 1.1**B) // 8 * 8 def n_layers(d): return int(math.ceil(d * 1.2**B)) def r(): return int(224 * 1.15**B) super(EfficientNet, self).__init__( # Stage 1 # nn.UpsamplingBilinear2d(size=(r(), r())), tu.kaiming(tnn.Conv3x3(in_ch, ch(32), stride=2, bias=False)), nn.BatchNorm2d(ch(32)), tnn.HardSwish(), # Stage 2 MBConv(ch(32), ch(16), 3, mul_factor=1), *[ MBConv(ch(16), ch(16), 3, mul_factor=1) for _ in range(n_layers(1) - 1) ], # Stage 3 MBConv(ch(16), ch(24), 3, stride=2), *[MBConv(ch(24), ch(24), 3) for _ in range(n_layers(2) - 1)], # Stage 4 MBConv(ch(24), ch(40), 5, stride=2), *[MBConv(ch(40), ch(40), 5) for _ in range(n_layers(2) - 1)], # Stage 5 MBConv(ch(40), ch(80), 3, stride=2), *[MBConv(ch(80), ch(80), 3) for _ in range(n_layers(3) - 1)], # Stage 6 MBConv(ch(80), ch(112), 5), *[MBConv(ch(112), ch(112), 5) for _ in range(n_layers(3) - 1)], # Stage 7 MBConv(ch(112), ch(192), 5, stride=2), *[MBConv(ch(192), ch(192), 5) for _ in range(n_layers(4) - 1)], # Stage 8 MBConv(ch(192), ch(320), 3), *[MBConv(ch(320), ch(320), 3) for _ in range(n_layers(1) - 1)], tu.kaiming(tnn.Conv1x1(ch(320), ch(1280), bias=False)), nn.BatchNorm2d(ch(1280)), tnn.HardSwish(), nn.AdaptiveAvgPool2d(1), tnn.Reshape(-1), tu.xavier(nn.Linear(ch(1280), num_classes)))
def VggGeneratorDebug(in_noise=32, out_ch=3): """ A not so small Vgg net image GAN generator for testing purposes Args: in_noise (int): dimension of the input noise out_ch (int): number of output channels (3 for RGB images) Returns: a VGG instance """ return nn.Sequential( kaiming(nn.Linear(in_noise, 128 * 16)), tnn.Reshape(128, 4, 4), nn.LeakyReLU(0.2, inplace=True), VggBNBone([128, 'U', 64, 'U', 32, 'U', 16], in_ch=128), xavier(tnn.Conv1x1(16, 1)), nn.Sigmoid())
def __init__(self, arch, n_skip_max=2, in_noise=256, out_ch=3): super(AutoGAN, self).__init__() self.n_skip_max = n_skip_max self.make_noise = (nn.Linear(in_noise, 4 * 4 * arch[0])) in_ch = arch[0] blocks = [] lasts = [] for i, out in enumerate(arch[1:]): mode = 'nearest' if i % 2 == 0 else 'bilinear' blocks.append(tnn.AutoGANGenBlock(in_ch, out, lasts, mode=mode)) lasts = ([out] + lasts)[:n_skip_max] in_ch = out self.blocks = nn.ModuleList(blocks) self.to_rgb = nn.Sequential(nn.BatchNorm2d(arch[-1]), nn.ReLU(True), xavier(tnn.Conv3x3(arch[-1], out_ch)))
def __init__(self, in_noise, out_ch, side_ch=1): super(VggImg2ImgGeneratorDebug, self).__init__() def make_block(in_ch, out_ch, **kwargs): return tnn.Conv2dNormReLU( in_ch, out_ch, norm=lambda out: tnn.Spade2d(out, side_ch, 64), **kwargs) self.net = tnn.CondSeq( kaiming(nn.Linear(in_noise, 128 * 16)), tnn.Reshape(128, 4, 4), nn.LeakyReLU(0.2, inplace=True), VggBNBone([128, 'U', 64, 'U', 32, 'U', 16], in_ch=128, block=make_block), xavier(tnn.Conv1x1(16, out_ch)), nn.Sigmoid())
def __init__(self, in_noise, out_ch, num_classes): super(VggClassCondGeneratorDebug, self).__init__() def make_block(in_ch, out_ch, **kwargs): return tnn.Conv2dNormReLU( in_ch, out_ch, norm=lambda out: tnn.ConditionalBN2d(out, 64), **kwargs) self.emb = nn.Embedding(num_classes, 64) self.net = tnn.CondSeq( kaiming(nn.Linear(in_noise, 128 * 16)), tnn.Reshape(128, 4, 4), nn.LeakyReLU(0.2, inplace=True), VggBNBone([128, 'U', 64, 'U', 32, 'U', 16], in_ch=128, block=make_block), xavier(tnn.Conv1x1(16, out_ch)), nn.Sigmoid())
def __init__(self, channels, cond_channels, hidden, momentum=0.8): super(Spade2d, self).__init__(channels, momentum) self.initial = kaiming(Conv3x3(cond_channels, hidden, stride=2)) self.make_weight = xavier(Conv3x3(hidden, channels)) self.make_bias = xavier(Conv3x3(hidden, channels)) self.register_buffer('weight', torch.ones(channels))