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)