def resnext(self) -> 'ResBlockBottleneck': self.wide(divisor=2) c = self.branch.conv2 assert isinstance(c, nn.Conv2d) self.branch.conv2 = kaiming( nn.Conv2d(c.in_channels, c.out_channels, kernel_size=3, padding=1, stride=self.stride, bias=c.bias is not None, groups=32)) return self
def resnext(self, groups: int = 32) -> 'PreactResBlockBottleneck': self.wide(divisor=4) c = self.branch.conv2 assert isinstance(c, nn.Conv2d) self.branch.conv2 = kaiming( nn.Conv2d(c.in_channels, c.out_channels, 3, groups=groups, stride=self.stride, padding=1, bias=False)) return self
def to_input_specs(self, in_channels: int) -> 'ConvBlock': """ Recreate a convolution with :code:`in_channels` input channels """ self.in_channels = in_channels c = self.conv self.conv = kaiming( Conv2d(in_channels, c.out_channels, c.kernel_size[0], stride=c.stride, bias=c.bias)) return self
def reset(self) -> None: """ Recreate the block as a simple conv-BatchNorm-ReLU """ self.add_module( 'conv', kaiming( Conv2d(self.in_channels, self.out_channels, self.kernel_size[0], stride=self.stride, bias=False))) self.add_module('norm', nn.BatchNorm2d(self.out_channels)) self.add_module('relu', nn.ReLU(inplace=True))
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 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, 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, 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 __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, in_ch, out_ch, cond_ch, ks, leak=0): super(Conv2dCondBNReLU, self).__init__() self.conv = kaiming(Conv2d(in_ch, out_ch, ks), a=leak) self.cbn = ConditionalBN2d(out_ch, cond_ch) self.leak = leak
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))
trainset = CachedDataset(trainset, transform=tfm) testset = CachedDataset(testset, transform=tfm_test) trainloader = DataLoader(trainset, args.batch_size, num_workers=8, pin_memory=True, shuffle=True, drop_last=True) testloader = DataLoader(testset, args.batch_size, num_workers=8, pin_memory=True) model = tvmodels.resnet18(pretrained=False) model.fc = tu.kaiming(torch.nn.Linear(512, len(testset.classes))) clf_recipe = CrossEntropyClassification(model, trainloader, testloader, testset.classes, log_every=10, test_every=50, lr=args.lr, beta1=args.beta1, wd=args.wd, visdom_env=args.visdom_env) clf_recipe.to(args.device) clf_recipe.run(args.epochs)
def to_equal_lr(self, leak=0.2) -> 'ResidualDiscriminator': for m in self.modules(): if isinstance(m, (nn.Linear, nn.Conv2d)): tu.kaiming(m, dynamic=True, a=leak) return self
def set_input_specs(self, in_channels: int) -> 'VGG': c1 = self.features.conv_1_1 assert isinstance(c1, tnn.ConvBlock) c1.conv = kaiming(tnn.Conv3x3(in_channels, c1.conv.out_channels, 3)) return self
def to_equal_lr(self) -> 'LinearReLU': if isinstance(self.relu, nn.LeakyReLU): tu.kaiming(self.linear, dynamic=True, a=self.relu.negative_slope) else: tu.kaiming(self.linear, dynamic=True) return self
def __init__(self, in_features: int, out_features) -> None: super().__init__() self.linear = tu.kaiming(nn.Linear(in_features, out_features)) self.relu = nn.ReLU(True)