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, ch: int, n_down: int, n_trunk: int = 2, n_post: int = 1, n_pre: int = 1, n_att_conv: int = 2, with_skips: bool = True) -> None: super(AttentionBlock, self).__init__() self.pre = tnn.CondSeq(*[Block(ch, ch) for _ in range(n_pre)]) self.post = tnn.CondSeq(*[Block(ch, ch) for _ in range(n_post)]) self.trunk = tnn.CondSeq(*[Block(ch, ch) for _ in range(n_trunk)]) soft: nn.Module = UBlock1(ch) for _ in range(n_down - 1): soft = UBlock(ch, soft, with_skip=with_skips) if n_down >= 0: conv1 = [soft] for i in range(n_att_conv): conv1 += [ nn.BatchNorm2d(ch), nn.ReLU(True), tu.kaiming(tnn.Conv1x1(ch, ch, bias=(i != n_att_conv - 1))) ] conv1.append(nn.Sigmoid()) self.mask = tnn.CondSeq(*conv1) else: self.mask = None
def __init__(self, ch, inner): super(UBlock, self).__init__() self.inner = inner if inner is not None: self.skip = Block(ch, ch) self.encode = tnn.CondSeq(nn.MaxPool2d(3, 2, 1), Block(ch, ch)) self.decode = tnn.CondSeq(Block(ch, ch), nn.UpsamplingBilinear2d(scale_factor=2))
def __init__(self, ch, n_down): super(AttentionBlock, self).__init__() self.pre = Block(ch, ch) self.post = Block(ch, ch) self.trunk = tnn.CondSeq(Block(ch, ch), Block(ch, ch)) soft = None for _ in range(n_down): soft = UBlock(ch, soft) self.mask = tnn.CondSeq(soft, nn.BatchNorm2d(ch), nn.ReLU(True), tu.kaiming(tnn.Conv1x1(ch, ch, bias=False)), nn.BatchNorm2d(ch), nn.ReLU(True), tu.kaiming(tnn.Conv1x1(ch, ch)), tnn.HardSigmoid())
def __init__(self, in_ch=3): super(Attention56Bone, self).__init__() self.head = tnn.CondSeq(tu.kaiming(tnn.Conv2d(in_ch, 64, 7, stride=2)), nn.ReLU(True), nn.MaxPool2d(3, 2, 1)) self.pre1 = Block(64, 256) self.attn1 = AttentionBlock(256, 3) self.pre2 = Block(256, 512, stride=2) self.attn2 = AttentionBlock(512, 2) self.pre3 = Block(512, 1024, stride=2) self.attn3 = AttentionBlock(1024, 1) self.pre4 = tnn.CondSeq( Block(1024, 2048, stride=2), Block(2048, 2048), Block(2048, 2048), )
def __init__(self, ch: int, inner: Optional[nn.Module], with_skip: bool = True) -> None: super(UBlock, self).__init__() self.inner = inner if with_skip and inner is not None: self.skip = Block(ch, ch) else: self.skip = None self.encode = tnn.CondSeq(nn.MaxPool2d(3, 1, 1), nn.UpsamplingBilinear2d(scale_factor=0.5), Block(ch, ch)) self.decode = tnn.CondSeq(Block(ch, ch), nn.UpsamplingBilinear2d(scale_factor=2))
def __init__(self, arch: list, num_classes: int) -> None: super().__init__() self.arch = arch in_ch = 3 self.in_channels = in_ch feats = tnn.CondSeq() block_num = 1 conv_num = 1 for layer in arch: if layer == 'M': feats.add_module(f'pool_{block_num}', nn.MaxPool2d(2, 2)) block_num += 1 conv_num = 1 else: ch = cast(int, layer) feats.add_module( f'conv_{block_num}_{conv_num}', tnn.ConvBlock(in_ch, ch, 3).remove_batchnorm()) in_ch = ch conv_num += 1 self.out_channels = ch self.features = feats self.classifier = ClassificationHead(self.out_channels, num_classes) self.classifier.to_vgg_style(4096)
def __init__(self, in_ch, hid, out_ch, quant_lvls, sz, n_layer=3): super(PixCNNBase, self).__init__() self.sz = sz self.lin = tnn.CondSeq( tnn.TopLeftConv2d(in_ch, hid, 5, center=False, bias=sz), nn.ReLU(inplace=True)) sz2 = sz[0] // 2, sz[1] // 2 sz4 = sz[0] // 4, sz[1] // 4 self.l1 = nn.Sequential( *[ResBlk(hid, hid * 2, hid, 5, sz) for _ in range(n_layer)]) self.l2 = nn.Sequential( *[ResBlk(hid, hid * 2, hid, 5, sz2) for _ in range(n_layer)]) self.l3 = nn.Sequential( *[ResBlk(hid, hid * 2, hid, 5, sz4) for _ in range(n_layer)]) self.l4 = nn.Sequential( *[ResBlk(hid, hid * 2, hid, 5, sz4) for _ in range(n_layer)]) self.l4 = nn.Sequential( *[ResBlk(hid, hid * 2, hid, 5, sz4) for _ in range(n_layer)]) self.l5 = nn.Sequential(*[ ResBlk(hid * 2, hid * 4, hid * 2, 5, sz2) for _ in range(n_layer) ]) self.l6 = nn.Sequential( * [ResBlk(hid * 3, hid * 6, hid * 3, 5, sz) for _ in range(n_layer)]) self.lout = PixelPredictor(hid * 3, out_ch)
def __init__(self, arch: List[Union[str, int]]) -> None: super().__init__() self.arch = arch self.in_channels = 3 in_ch = arch[0] features = tnn.CondSeq() assert isinstance(in_ch, int) features.add_module('input', tnn.Conv3x3(3, in_ch)) ii = 0 for i, (x, x2) in enumerate(zip(arch, arch[1:] + ['dummy'])): if x == 'D': continue downsample = x2 == 'D' assert isinstance(x, int) features.add_module(f'block_{ii}', tnn.ResidualDiscrBlock(in_ch, x, downsample)) in_ch = x ii += 1 self.out_channels = in_ch features.add_module('final_relu', nn.LeakyReLU(0.2, True)) assert isinstance(features.block_0, tnn.ResidualDiscrBlock) features.block_0.preact_skip() self.features = features self.classifier = ClassificationHead(self.out_channels, 1)
def to_standard_arch(self): self._modules.clear() arch = self.arch self.input = tnn.ConvBlock(3, int(arch[0]), 7) ch, i = int(arch[0]), 1 ii = 0 self.encode = tnn.CondSeq() while arch[i][0] == 'd': out_ch = int(arch[i][1:]) self.encode.add_module(f'conv_{ii}', tnn.ConvBlock(ch, out_ch, 3, stride=2)) ch = out_ch i += 1 ii += 1 ii = 0 self.transform = tnn.CondSeq() while arch[i][0] == 'R': out_ch = int(arch[i][1:]) self.transform.add_module(f'transform_{ii}', tnn.PreactResBlock(ch, out_ch)) ch = out_ch i += 1 ii += 1 ii = 0 self.decode = tnn.CondSeq() while i < len(arch) and arch[i][0] == 'u': out_ch = int(arch[i][1:]) self.decode.add_module( f'out_conv_{ii}', tnn.ConvBlock(ch, out_ch, 3, stride=1).add_upsampling()) ch = out_ch i += 1 ii += 1 self.to_rgb = tnn.ConvBlock(out_ch, 3, 7).remove_batchnorm() self.to_rgb.relu = nn.Sigmoid() def to_instance_norm(m): if isinstance(m, nn.BatchNorm2d): return nn.InstanceNorm2d(m.num_features, affine=True) if isinstance(m, nn.Conv2d): m.padding_mode = 'reflect' return m tnn.utils.edit_model(self, to_instance_norm)
def __init__(self, num_classes: int) -> None: super(Attention56Bone, self).__init__( OrderedDict([ ('head', tnn.CondSeq(tu.kaiming(tnn.Conv2d(3, 64, 7, stride=2)), nn.ReLU(True), nn.MaxPool2d(3, 2, 1))), ('pre1', Block(64, 256)), ('attn1', AttentionBlock(256, 3)), ('pre2', Block(256, 512, stride=2)), ('attn2', AttentionBlock(512, 2)), ('pre3', Block(512, 1024, stride=2)), ('attn3', AttentionBlock(1024, 1)), ('pre4', tnn.CondSeq( Block(1024, 2048, stride=2), Block(2048, 2048), Block(2048, 2048), )), ('classifier', ClassificationHead(2048, num_classes)) ]))
def _parse_snres(arch, in_ch): blocks = [] for x, x2 in zip(arch, arch[1:] + ['dummy']): if x == 'D': continue downsample = x2 == 'D' blocks.append(tnn.SNResidualDiscrBlock(in_ch, x, downsample)) in_ch = x return tnn.CondSeq(*blocks), in_ch
def __init__(self, in_ch, hid_ch, out_ch, ks, sz): super(ResBlk, self).__init__() self.go = tnn.CondSeq( nn.BatchNorm2d(in_ch), nn.ReLU(inplace=False), tnn.Conv1x1(in_ch, hid_ch), nn.BatchNorm2d(hid_ch), nn.ReLU(inplace=True), tnn.TopLeftConv2d(hid_ch, hid_ch, ks, center=True, bias=sz), nn.BatchNorm2d(hid_ch), nn.ReLU(inplace=True), tnn.Conv1x1(hid_ch, out_ch), )
def __init__(self, arch: List[str], num_classes: int) -> None: super().__init__() def parse(layer: str) -> List[int]: return [int(x) for x in layer.split(':')] self.arch = list(map(parse, arch)) self.features = tnn.CondSeq() self.features.add_module('input', ResNetInput(3, self.arch[0][0])) self._change_block_type('basic') self.classifier = ClassificationHead(self.arch[-1][0], num_classes)
def __init__(self, arch: List[int], num_classes: int) -> None: super().__init__() self.arch = arch self.in_channels = 3 self.out_channels = arch[-1] feats = tnn.CondSeq() feats.input = tnn.ConvBlock(3, arch[0], 3) encdec: nn.Module = tnn.ConvBlock(arch[-1], arch[-1] * 2, 3) for outer, inner in zip(arch[-2::-1], arch[:0:-1]): encdec = tnn.UBlock(outer, inner, encdec) feats.encoder_decoder = encdec self.features = feats assert isinstance(encdec.out_channels, int) self.classifier = tnn.ConvBlock(encdec.out_channels, num_classes, 3).remove_batchnorm().no_relu()
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, arch: List[int]) -> None: super().__init__() layers: List[nn.Module] = [ tnn.ConvBlock(3, arch[0], kernel_size=4, stride=2).remove_batchnorm().leaky() ] in_ch = arch[0] self.in_channels = in_ch for next_ch in arch[1:]: layers.append( tnn.ConvBlock(in_ch, next_ch, kernel_size=4, stride=2).leaky()) in_ch = next_ch assert isinstance(layers[-1], tnn.ConvBlock) layers[-1].conv.stride = (1, 1) self.features = tnn.CondSeq(*layers) self.classifier = tnn.Conv2d(in_ch, 1, 4)
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 VggBNBone(arch, in_ch=3, leak=0, block=tnn.Conv2dBNReLU, debug=False): """ Construct a VGG net How to specify a VGG architecture: It's a list of blocks specifications. Blocks are either: - 'M' for maxpool of kernel size 2 and stride 2 - 'A' for average pool of kernel size 2 and stride 2 - 'U' for nearest neighbors upsampling (scale factor 2) - an integer `ch` for a block with `ch` output channels Args: arch (list): architecture specification in_ch (int): number of input channels leak (float): leak in relus block (fn): block ctor Returns: A VGG instance """ layers = [] if debug: layers.append(tnn.Debug('Input')) for i, layer in enumerate(arch): if layer == 'M': layers.append(nn.MaxPool2d(2, 2)) elif layer == 'A': layers.append(nn.AvgPool2d(2, 2)) elif layer == 'U': layers.append(nn.UpsamplingNearest2d(scale_factor=2)) else: layers.append(block(in_ch, layer, ks=3, leak=leak)) in_ch = layer if debug: layer_name = 'layer_{}_{}'.format(layers[-1].__class__.__name__, i) layers.append(tnn.Debug(layer_name)) return tnn.CondSeq(*layers)
def _build(i, prev_ch): ch = int(arch[i][1:]) if arch[i][0] == 'R': transforms = tnn.CondSeq() transforms.in_channels = prev_ch ii = 0 while arch[i][0] == 'R': ch = int(arch[i][1:]) transforms.add_module(f'transform_{ii}', tnn.PreactResBlock(prev_ch, ch)) prev_ch = ch i += 1 ii += 1 transforms.out_channels = ch return transforms if arch[i][0] == 'd': u = tnn.encdec.UBlock(prev_ch, ch, _build(i + 1, ch)) u.to_bilinear_sampling() u.set_decoder_num_layers(1).set_encoder_num_layers(1) return u
def _change_block_type(self, ty: str) -> None: arch = self.arch feats = tnn.CondSeq() assert isinstance(self.features.input, ResNetInput) feats.add_module('input', self.features.input) in_ch = arch[0][0] self.in_channels = in_ch for i, (ch, s) in enumerate(arch): feats.add_module(f'block_{i}', self._make_block(ty, in_ch, ch, stride=s)) in_ch = ch self.out_channels = ch if 'preact' in ty: assert isinstance(feats.block_0, PREACT_BLOCKS) feats.block_0.no_preact() feats.add_module('final_bn', nn.BatchNorm2d(self.out_channels)) feats.add_module('final_relu', nn.ReLU(True)) self.features = feats
def ResNetBone(arch, head, block, in_ch=3, debug=False): """ A resnet How to specify an architecture: It's a list of block specifications. Each element is a string of the form "output channels:stride". For instance "64:2" is a block with input stride 2 and 64 output channels. Args: arch (list): the architecture specification head (fn): the module ctor to build for the first conv block (fn): the residual block to use ctor in_ch (int): number of input channels, 3 for RGB images debug (bool): should insert debug layers between each layer Returns: A Resnet instance """ def parse(l): return [int(x) for x in l.split(':')] layers = [] if debug: layers.append(tnn.Debug('Input')) ch, s = parse(arch[0]) layers.append(head(in_ch, ch, 7, stride=s)) if debug: layers.append(tnn.Debug('Head')) in_ch = ch for i, (ch, s) in enumerate(map(parse, arch[1:])): layers.append(block(in_ch, ch, stride=s)) in_ch = ch if debug: layer_name = 'layer_{}_{}'.format(layers[-1].__class__.__name__, i) layers.append(tnn.Debug(layer_name)) return tnn.CondSeq(*layers)
def base_patch_discr(arch, in_ch=3, out_ch=1, norm=None): def block(in_ch, out_ch, norm): if norm is None: return [ kaiming(nn.Conv2d(in_ch, out_ch, 4, stride=2, padding=1), a=0.2), nn.LeakyReLU(0.2, inplace=True) ] else: return [ kaiming(nn.Conv2d(in_ch, out_ch, 4, stride=2, padding=1), a=0.2), norm(out_ch), nn.LeakyReLU(0.2, inplace=True) ] layers = block(in_ch, arch[0], None) in_ch = arch[0] for out_ch in arch[1:]: layers += block(in_ch, out_ch, norm) in_ch = out_ch return tnn.CondSeq(*layers)
def __init__(self, ch): super(UBlock1, self).__init__() self.inner = tnn.CondSeq(nn.MaxPool2d(3, 1, 1), nn.UpsamplingBilinear2d(scale_factor=0.5), Block(ch, ch), nn.UpsamplingBilinear2d(scale_factor=2))