def __init__(self, nc=3, ndf=96):
super(NetD, self).__init__()
self.discriminator = nn.Sequential(
SpectralNorm(nn.Conv2d(nc, ndf, kernel_size=4, stride=2, padding=1, bias=False)), #bx1x64x64 --> bx96x32x32
nn.LeakyReLU(0.2, True),
SpectralNorm(nn.Conv2d(ndf, ndf*2, kernel_size=4, stride=2, padding=1, bias=False)), #bx96x32x32 --> bx192x16x16
# nn.BatchNorm2d(ndf*2),
nn.LeakyReLU(0.2, True),
SpectralNorm(nn.Conv2d(ndf*2, ndf*4, kernel_size=4, stride=2, padding=1, bias=False)), #bx192x16x16 --> bx384x8x8
# nn.BatchNorm2d(ndf*4),
nn.LeakyReLU(0.2, True),
SpectralNorm(nn.Conv2d(ndf*4, ndf*8, kernel_size=4, stride=2, padding=1, bias=False)),#bx384x8x8 --> bx768x4x4
# nn.BatchNorm2d(ndf*8),
nn.LeakyReLU(0.2, True),
)
self.attn = Self_Attn( ndf*8,'relu')
self.last = nn.Sequential(
nn.Conv2d(ndf*8, 1, kernel_size=4, stride=4, bias=False), #bx768x4x4--> bx1x1x1
# nn.Conv2d(ndf*8, 1, kernel_size=(10,3), stride=(10,3), bias=False),
# nn.Sigmoid() # for WGAN
)
def __init__(self, n_class=1000, chn=96, debug=False):
super().__init__()
def conv(in_channel, out_channel, downsample=True):
return GBlock(in_channel, out_channel,
bn=False,
upsample=False, downsample=downsample)
gain = 2 ** 0.5
if debug:
chn = 8
self.debug = debug
self.pre_conv = nn.Sequential(SpectralNorm(nn.Conv2d(3, 1*chn, 3,padding=1),),
nn.ReLU(),
SpectralNorm(nn.Conv2d(1*chn, 1*chn, 3,padding=1),),
nn.AvgPool2d(2))
self.pre_skip = SpectralNorm(nn.Conv2d(3, 1*chn, 1))
self.conv = nn.Sequential(conv(1*chn, 1*chn, downsample=True),
SelfAttention(1*chn),
conv(1*chn, 2*chn, downsample=True),
conv(2*chn, 4*chn, downsample=True),
conv(4*chn, 8*chn, downsample=True),
conv(8*chn, 16*chn, downsample=True),
conv(16*chn, 16*chn, downsample=False))
self.linear = SpectralNorm(nn.Linear(16*chn, 1))
self.embed = nn.Embedding(n_class, 16*chn)
self.embed.weight.data.uniform_(-0.1, 0.1)
self.embed = spectral_norm(self.embed)
def __init__(self, batch_size, image_size=64, z_dim=100, conv_dim=64, attn_feat=[16, 32], upsample=False):
super(Generator, self).__init__()
self.imsize = image_size
layers = []
n_layers = int(np.log2(self.imsize)) - 2
mult = 8 # 2 ** repeat_num # 8
assert mult * conv_dim > 3 * (2 ** n_layers), 'Need to add higher conv_dim, too many layers'
curr_dim = conv_dim * mult
# Initialize the first layer because it is different than the others.
layers.append(SpectralNorm(nn.ConvTranspose2d(z_dim, curr_dim, 4)))
layers.append(nn.BatchNorm2d(curr_dim))
layers.append(nn.ReLU())
for n in range(n_layers - 1):
layers.append(SpectralNorm(nn.ConvTranspose2d(curr_dim, int(curr_dim / 2), 4, 2, 1)))
layers.append(nn.BatchNorm2d(int(curr_dim / 2)))
layers.append(nn.ReLU())
# check the size of the feature space and add attention. (n+2) is used for indexing purposes
if 2**(n+2) in attn_feat:
layers.append(Self_Attn(int(curr_dim / 2), 'relu'))
curr_dim = int(curr_dim / 2)
# append a final layer to change to 3 channels and add Tanh activation
layers.append(nn.ConvTranspose2d(curr_dim, 3, 4, 2, 1))
layers.append(nn.Tanh())
self.output = nn.Sequential(*layers)
def __init__(self, channels):
super(ResidualBlock, self).__init__()
self.conv1 = SpectralNorm(nn.Conv2d(channels, channels, kernel_size=3, padding=1))
self.bn1 = nn.BatchNorm2d(channels)
self.prelu = nn.PReLU()
self.conv2 = SpectralNorm(nn.Conv2d(channels, channels, kernel_size=3, padding=1))
self.bn2 = nn.BatchNorm2d(channels)
def __init__(self, in_channel, out_channel, kernel_size=[3, 3],
padding=1, stride=1, n_class=None, bn=True,
activation=F.relu, upsample=True, downsample=False):
super().__init__()
gain = 2 ** 0.5
self.conv0 = SpectralNorm(nn.Conv2d(in_channel, out_channel,
kernel_size, stride, padding,
bias=True if bn else True))
self.conv1 = SpectralNorm(nn.Conv2d(out_channel, out_channel,
kernel_size, stride, padding,
bias=True if bn else True))
self.skip_proj = False
if in_channel != out_channel or upsample or downsample:
self.conv_sc = SpectralNorm(nn.Conv2d(in_channel, out_channel,
1, 1, 0))
self.skip_proj = True
self.upsample = upsample
self.downsample = downsample
self.activation = activation
self.bn = bn
if bn:
self.HyperBN = ConditionalNorm(in_channel, 148)
self.HyperBN_1 = ConditionalNorm(out_channel, 148)
def __init__(self, im_chan=10, conv_dim=64, leaky_relu_slope=0.2):
super(Discriminator, self).__init__()
self.l1 = nn.Sequential(
SpectralNorm(nn.Conv2d(im_chan, conv_dim, 4, 2, 1)),
nn.LeakyReLU(leaky_relu_slope),
)
curr_dim = conv_dim
self.l2 = nn.Sequential(
SpectralNorm(nn.Conv2d(curr_dim, curr_dim * 2, 4, 2, 1)),
nn.BatchNorm2d(curr_dim * 2),
nn.LeakyReLU(leaky_relu_slope),
)
curr_dim = curr_dim * 2
self.l3 = nn.Sequential(
SpectralNorm(nn.Conv2d(curr_dim, curr_dim * 2, 4, 2, 1)),
nn.BatchNorm2d(curr_dim * 2),
nn.LeakyReLU(leaky_relu_slope),
)
curr_dim = curr_dim * 2
self.l4 = nn.Sequential(
SpectralNorm(nn.Conv2d(curr_dim, curr_dim * 2, 4, 2, 1)),
nn.BatchNorm2d(curr_dim * 2),
nn.LeakyReLU(leaky_relu_slope),
)
curr_dim = curr_dim * 2
self.last = nn.Sequential(nn.Conv2d(curr_dim, 1, 4), )
self.attn1 = Self_Attn(256, 'relu')
self.attn2 = Self_Attn(512, 'relu')
def __init__(self, batch_size=64, image_size=64, conv_dim=64, rgb_channel=3):
super(Discriminator_DC, self).__init__()
self.imsize = image_size
layer1 = []
layer2 = []
layer3 = []
last = []
layer1.append(SpectralNorm(nn.Conv2d(rgb_channel, conv_dim, 4, 2, 1)))
layer1.append(nn.LeakyReLU(0.1))
curr_dim = conv_dim
layer2.append(SpectralNorm(nn.Conv2d(curr_dim, curr_dim * 2, 4, 2, 1)))
layer2.append(nn.LeakyReLU(0.1))
curr_dim = curr_dim * 2
layer3.append(SpectralNorm(nn.Conv2d(curr_dim, curr_dim * 2, 4, 2, 1)))
layer3.append(nn.LeakyReLU(0.1))
curr_dim = curr_dim * 2
if self.imsize == 64:
layer4 = []
layer4.append(SpectralNorm(nn.Conv2d(curr_dim, curr_dim * 2, 4, 2, 1)))
layer4.append(nn.LeakyReLU(0.1))
self.l4 = nn.Sequential(*layer4)
curr_dim = curr_dim * 2
self.l1 = nn.Sequential(*layer1)
self.l2 = nn.Sequential(*layer2)
self.l3 = nn.Sequential(*layer3)
last.append(nn.Conv2d(curr_dim, 1, 4))
self.last = nn.Sequential(*last)
def __init__(self,
batch_size=64,
image_size=64,
conv_dim=64,
attn_feat=[16, 32]):
super(Discriminator, self).__init__()
self.imsize = image_size
layers = []
n_layers = int(np.log2(self.imsize)) - 2
# Initialize the first layer because it is different than the others.
layers.append(SpectralNorm(nn.Conv2d(3, conv_dim, 4, 2, 1)))
layers.append(nn.LeakyReLU(0.1))
curr_dim = conv_dim
for n in range(n_layers - 1):
layers.append(
SpectralNorm(nn.Conv2d(curr_dim, curr_dim * 2, 4, 2, 1)))
layers.append(nn.LeakyReLU(0.1))
curr_dim *= 2
if 2**(n + 2) in attn_feat:
layers.append(Self_Attn_dynamic(curr_dim, 'relu'))
layers.append(nn.Conv2d(curr_dim, 1, 4))
self.output = nn.Sequential(*layers)
def __init__(self, batch_size=64, attn=True, image_size=28, conv_dim=64):
super().__init__()
self.imsize = image_size
layer1 = []
layer1.append(SpectralNorm(nn.Conv2d(1, conv_dim, 4, 2, 1)))
layer1.append(nn.LeakyReLU(0.1))
curr_dim = conv_dim
self.l1 = nn.Sequential(*layer1)
layer2 = []
layer2.append(SpectralNorm(nn.Conv2d(curr_dim, curr_dim * 2, 4, 2, 1)))
layer2.append(nn.LeakyReLU(0.1))
curr_dim = curr_dim * 2
self.l2 = nn.Sequential(*layer2)
layer3 = []
layer3.append(SpectralNorm(nn.Conv2d(curr_dim, curr_dim * 2, 4, 2, 1)))
layer3.append(nn.LeakyReLU(0.1))
curr_dim = curr_dim * 2
self.l3 = nn.Sequential(*layer3)
self.attn = Self_Attn(curr_dim, 'relu')
last = []
last.append(nn.Conv2d(curr_dim, 1, 4, 2, 1))
self.last = nn.Sequential(*last)
def __init__(self, batch_size=64, attn=True, image_size=28, z_dim=100, conv_dim=64):
super().__init__()
self.imsize = image_size
# Layer 1 turn 100 dims -> 512 dims, size 1 -> 3
layer1 = []
layer1.append(SpectralNorm(nn.ConvTranspose2d(in_channels = z_dim, out_channels = conv_dim*8, kernel_size = 3)))
layer1.append(nn.BatchNorm2d(conv_dim*8))
layer1.append(nn.ReLU())
self.l1 = nn.Sequential(*layer1)
# Layer 2 turn 512 dims -> 256 dims, size 3 -> 7
layer2 = []
layer2.append(SpectralNorm(nn.ConvTranspose2d(in_channels = conv_dim*8, out_channels = conv_dim*4,
kernel_size = 3, stride = 2, padding = 0)))
layer2.append(nn.BatchNorm2d(conv_dim*4))
layer2.append(nn.ReLU())
self.l2 = nn.Sequential(*layer2)
# Layer 3 turn 256 dims -> 128 dims, size 3 -> 14
layer3 = []
layer3.append(SpectralNorm(nn.ConvTranspose2d(in_channels = conv_dim*4, out_channels = conv_dim*2,
kernel_size = 4, stride = 2, padding = 1)))
layer3.append(nn.BatchNorm2d(conv_dim*2))
layer3.append(nn.ReLU())
self.l3 = nn.Sequential(*layer3)
# Layer 4 (Attn) turn 128 dims -> 128 dims
self.attn = Self_Attn(conv_dim*2, 'relu')
# Layer 5 turn 128 dims -> 1 dims, size 14 -> 28
last = []
last.append(nn.ConvTranspose2d(conv_dim*2, 1, 4, 2, 1))
last.append(nn.Tanh())
self.last = nn.Sequential(*last)
def __init__(self, batch_size, image_size=64, c_dim=5, conv_dim=64):
super(Generator, self).__init__()
self.imsize = image_size
layer1 = []
# 3x64x64 -> 64x32x32
layer1.append(SpectralNorm(nn.Conv2d(3 + c_dim, conv_dim, 4, 2, 1)))
layer1.append(nn.BatchNorm2d(conv_dim))
layer1.append(nn.ReLU(inplace=True))
# Down-sampling layers.
curr_dim = conv_dim
# 64x32x32 -> 128x16x16
layer1.append(SpectralNorm(nn.Conv2d(curr_dim, curr_dim * 2, 4, 2, 1)))
layer1.append(nn.BatchNorm2d(curr_dim * 2))
layer1.append(nn.ReLU(inplace=True))
curr_dim = curr_dim * 2
self.l1 = nn.Sequential(*layer1)
# 128x16x16
self.attn1 = Self_Attn(128, 'relu') # 256
# 128x16x16 -> 256x8x8
layer2 = []
layer2.append(SpectralNorm(nn.Conv2d(curr_dim, curr_dim * 2, 4, 2, 1)))
layer2.append(nn.BatchNorm2d(curr_dim * 2))
layer2.append(nn.ReLU(inplace=True))
curr_dim = curr_dim * 2
# Bottleneck layers. For testing whether sa-stargan outperforms stargan
#repeat_num = 1 #int(np.log2(self.imsize)) - 3
#for _ in range(1): # 3 for imsize=image_size=64
# layers.append(ResidualBlock(dim_in=curr_dim, dim_out=curr_dim))
# Up-sampling layers.
# 256x8x8 -> 128x16x16
layer2.append(
SpectralNorm(nn.ConvTranspose2d(curr_dim, curr_dim // 2, 4, 2, 1)))
layer2.append(nn.BatchNorm2d(curr_dim // 2))
layer2.append(nn.ReLU())
curr_dim = curr_dim // 2
self.l2 = nn.Sequential(*layer2)
# 128x16x16
self.attn2 = Self_Attn(128, 'relu')
# 128x16x16 -> 64x32x32
layer3 = []
layer3.append(
SpectralNorm(nn.ConvTranspose2d(curr_dim, curr_dim // 2, 4, 2, 1)))
layer3.append(nn.BatchNorm2d(curr_dim // 2))
layer3.append(nn.ReLU())
curr_dim = curr_dim // 2
layer3.append(nn.ConvTranspose2d(curr_dim, 3, 4, 2, 1))
layer3.append(nn.Tanh())
self.l3 = nn.Sequential(*layer3)
def __init__(self,
batch_size,
image_size=64,
z_dim=100,
conv_dim=64,
rgb_channel=3):
super(Generator_INV, self).__init__()
self.imsize = image_size
layer1 = []
layer2 = []
layer3 = []
last = []
repeat_num = int(np.log2(self.imsize)) - 3
mult = 2**repeat_num # 8
layer1.append(
SpectralNorm(nn.ConvTranspose2d(z_dim, conv_dim * mult, 4)))
layer1.append(nn.BatchNorm2d(conv_dim * mult))
layer1.append(nn.ReLU())
curr_dim = conv_dim * mult # =512
layer2.append(
SpectralNorm(
nn.ConvTranspose2d(curr_dim, int(curr_dim / 2), 4, 2, 1)))
layer2.append(nn.BatchNorm2d(int(curr_dim / 2)))
layer2.append(nn.ReLU())
curr_dim = int(curr_dim / 2)
layer3.append(
SpectralNorm(
nn.ConvTranspose2d(curr_dim, int(curr_dim / 2), 4, 2, 1)))
layer3.append(nn.BatchNorm2d(int(curr_dim / 2)))
layer3.append(nn.ReLU())
if self.imsize == 64:
layer4 = []
curr_dim = int(curr_dim / 2)
layer4.append(
SpectralNorm(
nn.ConvTranspose2d(curr_dim, int(curr_dim / 2), 4, 2, 1)))
layer4.append(nn.BatchNorm2d(int(curr_dim / 2)))
layer4.append(nn.ReLU())
self.l4 = nn.Sequential(*layer4)
curr_dim = int(curr_dim / 2)
self.l1 = nn.Sequential(*layer1)
self.l2 = nn.Sequential(*layer2)
self.l3 = nn.Sequential(*layer3)
last.append(nn.ConvTranspose2d(curr_dim, rgb_channel, 4, 2, 1))
last.append(nn.Tanh())
self.last = nn.Sequential(*last)
self.inv1 = Involution2d(64, 64, reduce_ratio=2)
self.inv2 = Involution2d(128, 128, reduce_ratio=2)
self.batchN2 = nn.BatchNorm2d(256)
def __init__(self,
batch_size,
image_size=64,
z_dim=100,
conv_dim=64,
rgb_channel=3):
super(Generator_SA, self).__init__()
self.imsize = image_size
layer1 = []
layer2 = []
layer3 = []
last = []
repeat_num = int(np.log2(self.imsize)) - 3
mult = 2**repeat_num # 8
layer1.append(
SpectralNorm(nn.ConvTranspose2d(z_dim, conv_dim * mult, 4)))
layer1.append(nn.BatchNorm2d(conv_dim * mult))
layer1.append(nn.ReLU())
curr_dim = conv_dim * mult
layer2.append(
SpectralNorm(
nn.ConvTranspose2d(curr_dim, int(curr_dim / 2), 4, 2, 1)))
layer2.append(nn.BatchNorm2d(int(curr_dim / 2)))
layer2.append(nn.ReLU())
curr_dim = int(curr_dim / 2)
layer3.append(
SpectralNorm(
nn.ConvTranspose2d(curr_dim, int(curr_dim / 2), 4, 2, 1)))
layer3.append(nn.BatchNorm2d(int(curr_dim / 2)))
layer3.append(nn.ReLU())
if self.imsize == 64:
layer4 = []
curr_dim = int(curr_dim / 2)
layer4.append(
SpectralNorm(
nn.ConvTranspose2d(curr_dim, int(curr_dim / 2), 4, 2, 1)))
layer4.append(nn.BatchNorm2d(int(curr_dim / 2)))
layer4.append(nn.ReLU())
self.l4 = nn.Sequential(*layer4)
curr_dim = int(curr_dim / 2)
self.l1 = nn.Sequential(*layer1)
self.l2 = nn.Sequential(*layer2)
self.l3 = nn.Sequential(*layer3)
last.append(nn.ConvTranspose2d(curr_dim, rgb_channel, 4, 2, 1))
last.append(nn.Tanh())
self.last = nn.Sequential(*last)
self.attn1 = Self_Attn(128, 'relu')
self.attn2 = Self_Attn(64, 'relu')
def __init__(self, batch_size, image_size=64, z_dim=100, conv_dim=64):
super(Generator, self).__init__()
self.imsize = image_size
layer1 = []
layer2 = []
layer3 = []
# layern = []
last = []
repeat_num = int(np.log2(self.imsize)) - 3
mult = 2 ** repeat_num # 8
layer1.append(SpectralNorm(nn.ConvTranspose2d(z_dim, conv_dim * mult, 4)))
layer1.append(nn.BatchNorm2d(conv_dim * mult))
layer1.append(nn.ReLU())
curr_dim = conv_dim * mult
layer2.append(SpectralNorm(nn.ConvTranspose2d(curr_dim, int(curr_dim / 2), 3, 2, 2))) # 4,2,1
layer2.append(nn.BatchNorm2d(int(curr_dim / 2)))
layer2.append(nn.ReLU())
curr_dim = int(curr_dim / 2)
layer3.append(SpectralNorm(nn.ConvTranspose2d(curr_dim, int(curr_dim / 2), 3, 2, 2)))
layer3.append(nn.BatchNorm2d(int(curr_dim / 2)))
layer3.append(nn.ReLU())
# curr_dim = int(curr_dim / 2)
#
# layern.append(SpectralNorm(nn.ConvTranspose1d(curr_dim, int(curr_dim / 2), 4, 2, 1)))
# layern.append(nn.BatchNorm2d(int(curr_dim / 2)))
# layern.append(nn.ReLU())
if self.imsize == 64:
layer4 = []
curr_dim = int(curr_dim / 2)
layer4.append(SpectralNorm(nn.ConvTranspose2d(curr_dim, int(curr_dim / 2), 4, 2, 1)))
layer4.append(nn.BatchNorm2d(int(curr_dim / 2)))
layer4.append(nn.ReLU())
self.l4 = nn.Sequential(*layer4)
curr_dim = int(curr_dim / 2)
# self.ln = nn.Sequential(*layern)
self.l1 = nn.Sequential(*layer1)
self.l2 = nn.Sequential(*layer2)
self.l3 = nn.Sequential(*layer3)
last.append(nn.ConvTranspose2d(64, 1, 2, 2, 1)) # curr_dim
last.append(nn.Tanh())
self.last = nn.Sequential(*last)
self.attn1 = Self_Attn( 64, 'relu') #128
self.attn2 = Self_Attn( 64, 'relu')
self.input1d2d = nn.ConvTranspose1d(144,128,1)#SpectralNorm(nn.ConvTranspose1d(144,128,1))
def __init__(self,
input_nc,
ndf=64,
n_layers=3,
use_sigmoid=False,
getIntermFeat=False):
super(NLayerDiscriminator, self).__init__()
self.getIntermFeat = getIntermFeat
self.n_layers = n_layers
kw = 4
padw = int(np.ceil((kw - 1.0) / 2))
sequence = [[
nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw),
nn.LeakyReLU(0.2, True)
]]
nf = ndf
for n in range(1, n_layers):
nf_prev = nf
nf = min(nf * 2, 512)
sequence += [[
SpectralNorm(
nn.Conv2d(nf_prev,
nf,
kernel_size=kw,
stride=2,
padding=padw)),
nn.LeakyReLU(0.2, True)
]]
nf_prev = nf
nf = min(nf * 2, 512)
sequence += [[
SpectralNorm(
nn.Conv2d(nf_prev, nf, kernel_size=kw, stride=1,
padding=padw)),
nn.LeakyReLU(0.2, True)
]]
sequence += [[
nn.Conv2d(nf, 1, kernel_size=kw, stride=1, padding=padw)
]]
if use_sigmoid:
sequence += [[nn.Sigmoid()]]
if getIntermFeat:
for n in range(len(sequence)):
setattr(self, 'model' + str(n), nn.Sequential(*sequence[n]))
else:
sequence_stream = []
for n in range(len(sequence)):
sequence_stream += sequence[n]
self.model = nn.Sequential(*sequence_stream)
def __init__(self):
super(Discriminator, self).__init__()
self.net = nn.Sequential(
SpectralNorm(nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1)),
# nn.BatchNorm2d(64),
nn.LeakyReLU(0.2),
SpectralNorm(nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1)),
# nn.BatchNorm2d(64),
nn.LeakyReLU(0.2),
SpectralNorm(nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1)),
# nn.BatchNorm2d(128),
nn.LeakyReLU(0.2),
SpectralNorm(nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1)),
# nn.BatchNorm2d(128),
nn.LeakyReLU(0.2),
SpectralNorm(nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1)),
# nn.BatchNorm2d(256),
nn.LeakyReLU(0.2),
SpectralNorm(nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)),
# nn.BatchNorm2d(256),
nn.LeakyReLU(0.2),
SpectralNorm(nn.Conv2d(256, 512, kernel_size=3, stride=2, padding=1)),
# nn.BatchNorm2d(512),
nn.LeakyReLU(0.2),
SpectralNorm(nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)),
# nn.BatchNorm2d(512),
nn.LeakyReLU(0.2),
# SpectralNorm(nn.Conv2d(512, 1, kernel_size=3, stride=1, padding=1)) # 1
# nn.AdaptiveAvgPool2d(1), # 2
# SpectralNorm(nn.Conv2d(512, 1, kernel_size=1, stride=1, padding=1)),
# nn.AdaptiveAvgPool2d(1), # 3
SpectralNorm(nn.Conv2d(512, 1024, kernel_size=1)),
nn.LeakyReLU(0.2),
SpectralNorm(nn.Conv2d(1024, 1, kernel_size=1))
)
def __init__(self, batch_size=64, image_size=64, z_dim=100, conv_dim=64):
super(Discriminator, self).__init__()
self.imsize = image_size
layer1 = []
layer2 = []
layer3 = []
last = []
## For inference x
layer1.append(SpectralNorm(nn.Conv2d(3, conv_dim, 4, 2, 1)))
layer1.append(nn.LeakyReLU(0.1))
curr_dim = conv_dim
layer2.append(SpectralNorm(nn.Conv2d(curr_dim, curr_dim * 2, 4, 2, 1)))
layer2.append(nn.LeakyReLU(0.1))
curr_dim = curr_dim * 2
layer3.append(SpectralNorm(nn.Conv2d(curr_dim, curr_dim * 2, 4, 2, 1)))
layer3.append(nn.LeakyReLU(0.1))
curr_dim = curr_dim * 2
if self.imsize == 64:
layer4 = []
layer4.append(SpectralNorm(nn.Conv2d(curr_dim, curr_dim * 2, 4, 2, 1)))
layer4.append(nn.LeakyReLU(0.1))
self.l4 = nn.Sequential(*layer4)
curr_dim = curr_dim*2
self.l1 = nn.Sequential(*layer1)
self.l2 = nn.Sequential(*layer2)
self.l3 = nn.Sequential(*layer3)
last.append(SpectralNorm(nn.Conv2d(curr_dim, curr_dim, 4)))
last.append(nn.LeakyReLU(0.1))
self.last = nn.Sequential(*last)
self.attn1 = Self_Attn(256, 'relu')
self.attn2 = Self_Attn(512, 'relu')
## For inference z
self.infer_z = nn.Sequential(
SpectralNorm(nn.Conv2d(z_dim, 512, 1)),
nn.LeakyReLU(0.1),
SpectralNorm(nn.Conv2d(512, 512, 1)),
nn.LeakyReLU(0.1),
)
## For inference joint x,z
self.infer_joint = nn.Sequential(
SpectralNorm(nn.Conv2d(1024, 1024, 1)),
nn.LeakyReLU(0.1),
SpectralNorm(nn.Conv2d(1024, 1024, 1)),
nn.LeakyReLU(0.1),
nn.Conv2d(1024, 1, 1),
nn.Sigmoid(),
)
def encode_image_by_16times(ndf):
layers = []
layers.append(SpectralNorm(nn.Conv2d(3, ndf, 4, 2, 1, bias=True)))
layers.append(nn.LeakyReLU(0.2, inplace=True), )
layers.append(SpectralNorm(nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=True)))
layers.append(nn.LeakyReLU(0.2, inplace=True))
layers.append(SpectralNorm(nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1,
bias=True)))
layers.append(nn.LeakyReLU(0.2, inplace=True))
layers.append(SpectralNorm(nn.Conv2d(ndf * 4, ndf * 8, 4, 2, 1,
bias=True)))
layers.append(nn.LeakyReLU(0.2, inplace=True))
return nn.Sequential(*layers)
def __init__(self,
code_dim=100,
n_class=1000,
chn=96,
img_size=128,
debug=False):
super().__init__()
self.linear = SpectralNorm(nn.Linear(n_class, 128, bias=False))
if debug:
chn = 8
# For 128x128 images
if img_size == 128:
self.first_view = 16 * chn
self.G_linear = SpectralNorm(nn.Linear(20, 4 * 4 * 16 * chn))
self.conv = nn.ModuleList([
GBlock(16 * chn, 16 * chn, n_class=n_class),
GBlock(16 * chn, 8 * chn, n_class=n_class),
GBlock(8 * chn, 4 * chn, n_class=n_class),
GBlock(4 * chn, 2 * chn, n_class=n_class),
SelfAttention(2 * chn),
GBlock(2 * chn, 1 * chn, n_class=n_class)
])
# For 512x512 images
elif img_size == 512:
self.first_view = 16 * chn
self.G_linear = SpectralNorm(nn.Linear(16, 4 * 4 * 16 * chn))
self.conv = nn.ModuleList([
GBlock(16 * chn, 16 * chn, n_class=n_class, n_condition=144),
GBlock(16 * chn, 8 * chn, n_class=n_class, n_condition=144),
GBlock(8 * chn, 8 * chn, n_class=n_class, n_condition=144),
GBlock(8 * chn, 4 * chn, n_class=n_class, n_condition=144),
SelfAttention(4 * chn),
GBlock(4 * chn, 2 * chn, n_class=n_class, n_condition=144),
GBlock(2 * chn, 1 * chn, n_class=n_class, n_condition=144),
GBlock(1 * chn, 1 * chn, n_class=n_class, n_condition=144)
])
# TODO impl ScaledCrossReplicaBatchNorm
self.ScaledCrossReplicaBN = ScaledCrossReplicaBatchNorm2d(1 * chn)
self.colorize = SpectralNorm(nn.Conv2d(1 * chn, 3, [3, 3], padding=1))
def __init__(self, in_nc, out_nc, nf):
super(generator, self).__init__()
self.input_nc = in_nc
self.output_nc = out_nc
self.nf = nf
self.downconv1 = nn.Sequential( # input H,W 3 output H,W 64 1
SpectralNorm(nn.Conv2d(in_nc, nf, 5, 1, 2)),
nn.InstanceNorm2d(nf),
nn.ReLU(True),
)
self.downconv2 = nn.Sequential( # input H,W 64 output H/2,W/2 128 2
SpectralNorm(nn.Conv2d(nf, nf * 2, 3, 2, 1)),
nn.InstanceNorm2d(nf * 2),
nn.ReLU(True),
)
self.downconv3 = nn.Sequential( # input H/2,W/2 128 output H/4,W/4 256 3
SpectralNorm(nn.Conv2d(nf * 2, nf * 4, 3, 2, 1)),
nn.InstanceNorm2d(nf * 4),
nn.ReLU(True),
)
self.downconv4 = nn.Sequential( # input H/4,W/4 256 output H/8,W/8 512 4
SpectralNorm(nn.Conv2d(nf * 4, nf * 8, 3, 2, 1)),
nn.InstanceNorm2d(nf * 8),
nn.ReLU(True),
)
self.downconv5 = nn.Sequential( # input H/8,W/8 512 output H/8,W/8 512 5
SpectralNorm(nn.Conv2d(nf * 8, nf * 8, 1, 1)),
nn.InstanceNorm2d(nf * 8),
nn.ReLU(True),
)
self.upconv3 = nn.Sequential( # input H/8,W/8 1024 output H/4,W/4 256 6
SpectralNorm(nn.ConvTranspose2d(nf * 16, nf * 4, 4, 2, 1)),
nn.InstanceNorm2d(nf * 4),
nn.ReLU(True),
)
self.upconv2 = nn.Sequential( # input H/4,W/4 512 output H/2,W/2 128 7
SpectralNorm(nn.ConvTranspose2d(nf * 8, nf * 2, 4, 2, 1)),
nn.InstanceNorm2d(nf * 2),
nn.ReLU(True),
)
self.upconv1 = nn.Sequential( # input H/2,W/2 256 output H,W 3 8
SpectralNorm(nn.ConvTranspose2d(nf * 4, nf, 4, 2, 1)),
nn.InstanceNorm2d(nf),
nn.ReLU(True),
)
self.CAM = ChannelAttentionModule()
def __init__(self, nc=3, ndf=96):
super(NetA, self).__init__()
self.discriminator1 = nn.Sequential(
SpectralNorm(
nn.Conv2d(nc * 2,
ndf,
kernel_size=4,
stride=2,
padding=1,
bias=False)),
nn.LeakyReLU(0.2, True),
SpectralNorm(
nn.Conv2d(ndf,
ndf * 2,
kernel_size=4,
stride=2,
padding=1,
bias=False)),
# nn.BatchNorm2d(ndf*2),
nn.LeakyReLU(0.2, True),
SpectralNorm(
nn.Conv2d(ndf * 2,
ndf * 4,
kernel_size=4,
stride=2,
padding=1,
bias=False)),
# nn.BatchNorm2d(ndf*4),
nn.LeakyReLU(0.2, True),
)
self.attn1 = Self_Attn(ndf * 4, 'relu')
self.discriminator2 = nn.Sequential(
SpectralNorm(
nn.Conv2d(ndf * 4,
ndf * 8,
kernel_size=4,
stride=2,
padding=1,
bias=False)),
# nn.BatchNorm2d(ndf*8),
nn.LeakyReLU(0.2, True),
)
self.attn2 = Self_Attn(ndf * 8, 'relu')
self.last = nn.Sequential(
nn.Conv2d(ndf * 8, 1, kernel_size=4, stride=4, bias=False),
# nn.Conv2d(ndf*8, 1, kernel_size=(10,3), stride=(10,3), bias=False),
# nn.Sigmoid() # for WGAN
)
def __init__(self, batch_size=64, image_size=128, conv_dim=64):
super(Discriminator, self).__init__()
self.imsize = image_size
layer1 = []
layer2 = []
layer3 = []
layer4 = []
layer5 = []
last = []
layer1.append(SpectralNorm(nn.Conv2d(3, conv_dim, 4, 2, 1))) #64, 3 64
layer1.append(nn.LeakyReLU(0.1))
curr_dim = conv_dim
layer2.append(SpectralNorm(nn.Conv2d(curr_dim, curr_dim * 2, 4, 2,
1))) #32, 64 128
layer2.append(nn.LeakyReLU(0.1))
curr_dim = curr_dim * 2
layer3.append(SpectralNorm(nn.Conv2d(curr_dim, curr_dim * 2, 4, 2,
1))) #16, 128 256
layer3.append(nn.LeakyReLU(0.1))
curr_dim = curr_dim * 2
layer4.append(SpectralNorm(nn.Conv2d(curr_dim, curr_dim * 2, 4, 2,
1))) #8, 256 512
layer4.append(nn.LeakyReLU(0.1))
curr_dim = curr_dim * 2
layer5.append(SpectralNorm(nn.Conv2d(curr_dim, curr_dim * 2, 4, 2,
1))) #4, 512 1024
layer5.append(nn.LeakyReLU(0.1))
curr_dim = curr_dim * 2
self.l1 = nn.Sequential(*layer1)
self.l2 = nn.Sequential(*layer2)
self.l3 = nn.Sequential(*layer3)
self.l4 = nn.Sequential(*layer4)
self.l5 = nn.Sequential(*layer5)
last.append(nn.Conv2d(curr_dim, 1, 4))
self.last = nn.Sequential(*last)
self.attn1 = Self_Attn(512, 'relu')
self.attn2 = Self_Attn(1024, 'relu')
def __init__(self, in_ch, out_ch, norm_layer, use_bias, use_spectral=False):
super(Up, self).__init__()
if use_spectral:
self.up = nn.Sequential(
# nn.Upsample(scale_factor=2, mode='nearest'),
# nn.Conv2d(in_ch, out_ch,
# kernel_size=3, stride=1,
# padding=1, bias=use_bias),
SpectralNorm(nn.ConvTranspose2d(in_ch, out_ch,
kernel_size=3, stride=2,
padding=1, output_padding=1,
bias=use_bias)),
norm_layer(out_ch),
nn.ReLU(True)
)
else:
self.up = nn.Sequential(
# nn.Upsample(scale_factor=2, mode='nearest'),
# nn.Conv2d(in_ch, out_ch,
# kernel_size=3, stride=1,
# padding=1, bias=use_bias),
nn.ConvTranspose2d(in_ch, out_ch,
kernel_size=3, stride=2,
padding=1, output_padding=1,
bias=use_bias),
norm_layer(out_ch),
nn.ReLU(True)
)
def __init__(self, batch_size, image_size=64, z_dim=100, conv_dim=64):
super(Generator, self).__init__()
self.imsize = image_size
self.layers = []
first_layer = []
loop_layers = []
attn_layers = []
final_layer = []
attn_feat = [16, 32]
n_layers = int(np.log2(self.imsize)) - 2
mult = 8 #2 ** repeat_num # 8
assert mult * conv_dim > 3 * (
2**n_layers), 'Need to add higher conv_dim, too many layers'
curr_dim = conv_dim * mult
# Initialize the first layer because it is different than the others.
first_layer.append(
SpectralNorm(nn.ConvTranspose2d(3, conv_dim, 4, 2, 1)))
self.layers.append(nn.BatchNorm2d(curr_dim))
first_layer.append(nn.ReLU())
for n in range(n_layers - 1):
loop_layers.append([])
loop_layers[-1].append(
SpectralNorm(
nn.ConvTranspose2d(curr_dim, int(curr_dim / 2), 4, 2, 1)))
loop_layers[-1].append(nn.BatchNorm2d(int(curr_dim / 2)))
loop_layers[-1].append(nn.ReLU())
if 2**(n + 2) in attn_feat:
attn_layers.append([])
attn_layers[-1].append(Self_Attn(int(curr_dim / 2), 'relu'))
curr_dim = int(curr_dim / 2)
# append a final layer to change to 3 channels and add Tanh activation
final_layer.append(nn.ConvTranspose2d(curr_dim, 3, 4, 2, 1))
final_layer.append(nn.Tanh())
self.layers.append(nn.Sequential(*first_layer))
for n in range(n_layers - 1):
self.layers.append(nn.Sequential(*loop_layers[n]))
if n == 1:
self.layers.append(attn_layers[0])
if n == 2:
self.layers.append(attn_layers[1])
self.layers.append(nn.Sequential(*final_layer))
def __init__(self, nc=3, ngf=96):
super(NetG, self).__init__()
self.encoder = nn.Sequential(
nn.Conv2d(nc, ngf, kernel_size=4, stride=2, padding=1, bias=False),#bx1x64x64 --> bx96x32x32
nn.LeakyReLU(0.2, True),
nn.Conv2d(ngf, ngf*2, kernel_size=4, stride=2, padding=1,#bx96x32x32 --> bx192x16x16
bias=False),
nn.BatchNorm2d(ngf*2),
nn.LeakyReLU(0.2, True),
nn.Conv2d(ngf*2, ngf*4, kernel_size=4, stride=2, padding=1, #bx192x16x16 --> bx384x8x8
bias=False),
nn.BatchNorm2d(ngf*4),
nn.LeakyReLU(0.2, True),
nn.Conv2d(ngf*4, ngf*8, kernel_size=4, stride=2, padding=1, #bx384x8x8 --> bx768x4x4
bias=False),
nn.BatchNorm2d(ngf*8),
nn.LeakyReLU(0.2, True),
)
self.decoder = nn.Sequential(
SpectralNorm(nn.ConvTranspose2d(ngf*8, ngf*4,
kernel_size=4, stride=2, padding=1, bias=False)),
nn.BatchNorm2d(ngf*4),
nn.ReLU(),
# nn.ReLU(True), #bug
SpectralNorm(nn.ConvTranspose2d(ngf*4, ngf*2, kernel_size=4,
stride=2, padding=1, bias=False)),
nn.BatchNorm2d(ngf*2),
nn.ReLU(),
# nn.ReLU(True), #bug
SpectralNorm(nn.ConvTranspose2d(ngf*2, ngf, kernel_size=4,
stride=2, padding=1, bias=False)),
nn.BatchNorm2d(ngf),
nn.ReLU(),
# nn.ReLU(True),#bug
)
self.attn = Self_Attn( ngf, 'relu')
self.last = nn.Sequential(
nn.ConvTranspose2d(ngf, nc, kernel_size=4, stride=2, padding=1,
bias=False),
nn.Tanh()
)
def __init__(self, dim_in, dim_out):
super(ResidualBlock, self).__init__()
self.main = nn.Sequential(
SpectralNorm(
nn.Conv2d(dim_in,
dim_out,
kernel_size=3,
stride=1,
padding=1,
bias=False)), nn.ReLU(inplace=True),
SpectralNorm(
nn.Conv2d(dim_out,
dim_out,
kernel_size=3,
stride=1,
padding=1,
bias=False)))
def __init__(self, in_dim, out_dim):
super(ConvBNReLU, self).__init__()
self.in_dim = in_dim
self.out_dim = out_dim
self.conv = SpectralNorm(
nn.Conv2d(in_dim, out_dim, kernel_size=3, stride=1, padding=1))
self.bn = nn.BatchNorm2d(out_dim)
self.relu = nn.ReLU(True)
def __init__(self, batch_size=64, attn=True, image_size=64, z_dim=100, conv_dim=64):
super().__init__()
self.attn = attn
# Layer 1 turn 100 dims -> 512 dims, size 1 -> 4
layer1 = []
layer1.append(SpectralNorm(nn.ConvTranspose2d(in_channels = z_dim, out_channels = conv_dim*8, kernel_size = 4)))
layer1.append(nn.BatchNorm2d(conv_dim*8))
layer1.append(nn.ReLU())
self.l1 = nn.Sequential(*layer1)
# Layer 2 turn 512 dims -> 256 dims, size 4 -> 8
layer2 = []
layer2.append(SpectralNorm(nn.ConvTranspose2d(in_channels = conv_dim*8, out_channels = conv_dim*4,
kernel_size = 4, stride = 2, padding = 1)))
layer2.append(nn.BatchNorm2d(conv_dim*4))
layer2.append(nn.ReLU())
self.l2 = nn.Sequential(*layer2)
# Layer 3 turn 256 dims -> 128 dims, size 8 -> 16
layer3 = []
layer3.append(SpectralNorm(nn.ConvTranspose2d(in_channels = conv_dim*4, out_channels = conv_dim*2,
kernel_size = 4, stride = 2, padding = 1)))
layer3.append(nn.BatchNorm2d(conv_dim*2))
layer3.append(nn.ReLU())
self.l3 = nn.Sequential(*layer3)
# Attn1 layer turn 128 dims -> 128 dims
self.attn1 = Self_Attn(conv_dim*2)
# Layer 4 turn 128 dims -> 64 dims, size 16 -> 32
layer4 = []
layer4.append(SpectralNorm(nn.ConvTranspose2d(in_channels = conv_dim*2, out_channels = conv_dim,
kernel_size = 4, stride = 2, padding = 1)))
layer4.append(nn.BatchNorm2d(conv_dim))
layer4.append(nn.ReLU())
self.l4 = nn.Sequential(*layer4)
# Attn2 layer turn 64 dims -> 64 dims
self.attn2 = Self_Attn(conv_dim)
# Layer 5 turn 64 dims -> 3 dims, size 32 -> 64
layer5 = []
layer5.append(nn.ConvTranspose2d(conv_dim, 3, 4, 2, 1))
layer5.append(nn.Tanh())
self.l5 = nn.Sequential(*layer5)
def __init__(self,
batch_size=64,
image_size=64,
conv_dim=64,
sn_type='sn'):
super(Discriminator, self).__init__()
self.imsize = image_size
layer1 = []
layer2 = []
layer3 = []
last = []
if sn_type == 'sn':
from spectral import SpectralNorm
elif sn_type == 'adasn':
from adaSN import SpectralNorm
layer1.append(SpectralNorm(nn.Conv2d(3, conv_dim, 4, 2, 1)))
layer1.append(nn.LeakyReLU(0.1))
curr_dim = conv_dim
layer2.append(SpectralNorm(nn.Conv2d(curr_dim, curr_dim * 2, 4, 2, 1)))
layer2.append(nn.LeakyReLU(0.1))
curr_dim = curr_dim * 2
layer3.append(SpectralNorm(nn.Conv2d(curr_dim, curr_dim * 2, 4, 2, 1)))
layer3.append(nn.LeakyReLU(0.1))
curr_dim = curr_dim * 2
if self.imsize == 64:
layer4 = []
layer4.append(
SpectralNorm(nn.Conv2d(curr_dim, curr_dim * 2, 4, 2, 1)))
layer4.append(nn.LeakyReLU(0.1))
self.l4 = nn.Sequential(*layer4)
curr_dim = curr_dim * 2
self.l1 = nn.Sequential(*layer1)
self.l2 = nn.Sequential(*layer2)
self.l3 = nn.Sequential(*layer3)
last.append(nn.Conv2d(curr_dim, 1, 4))
self.last = nn.Sequential(*last)
self.attn1 = Self_Attn(256, 'relu')
self.attn2 = Self_Attn(512, 'relu')
def __init__(self):
super(Generator, self).__init__()
self.init_size = opt.img_size // 4
self.l1 = nn.Sequential(
nn.Linear(opt.latent_dim, 128 * self.init_size**2))
self.conv_blocks = nn.Sequential(
nn.BatchNorm2d(128), nn.Upsample(scale_factor=2),
SpectralNorm(nn.Conv2d(128, 128, 3, stride=1, padding=1)),
nn.BatchNorm2d(128, 0.8), nn.LeakyReLU(0.2, inplace=True),
nn.Upsample(scale_factor=2),
SpectralNorm(nn.Conv2d(128, 64, 3, stride=1, padding=1)),
nn.BatchNorm2d(64, 0.8), nn.LeakyReLU(0.2, inplace=True),
AttentionLayer(input_size=64, attention_size=64),
SpectralNorm(nn.Conv2d(64, opt.channels, 3, stride=1, padding=1)),
nn.Tanh())
