def __init__(self, enc_dim=64, enc_layers=5, enc_norm_fn='batchnorm', enc_acti_fn='lrelu',
dec_dim=64, dec_layers=5, dec_norm_fn='batchnorm', dec_acti_fn='relu',
n_attrs=13, shortcut_layers=0, inject_layers=0, img_size=128):
super(Generator, self).__init__()
self.shortcut_layers = min(shortcut_layers, dec_layers - 1)
self.inject_layers = min(inject_layers, dec_layers - 1)
self.f_size = img_size // 2**enc_layers # f_size = 4 for 128x128
layers = []
n_in = 3
# start
layers += [Conv2dBlock(3, 64, (4, 4), stride=2, padding=1, norm_fn=enc_norm_fn, acti_fn=enc_acti_fn)] # 0
layers += [Self_Attn(64, 'relu')] # 1
layers += [Conv2dBlock(64, 128, (4, 4), stride=2, padding=1, norm_fn=enc_norm_fn, acti_fn=enc_acti_fn)] # 2
layers += [Self_Attn(128, 'relu')] # 3
layers += [Conv2dBlock(128, 256, (4, 4), stride=2, padding=1, norm_fn=enc_norm_fn, acti_fn=enc_acti_fn)] # 4
layers += [Conv2dBlock(256, 512, (4, 4), stride=2, padding=1, norm_fn=enc_norm_fn, acti_fn=enc_acti_fn)] # 5
layers += [Conv2dBlock(512, 1024, (4, 4), stride=2, padding=1, norm_fn=enc_norm_fn, acti_fn=enc_acti_fn)] # 6
self.enc_layers = nn.ModuleList(layers)
# end
layers = []
n_in = n_in + n_attrs # 1024 + 13
k,j=0,0
if self.inject_layers == 1:
k=13
elif self.inject_layers == 2:
k=13
j=13
else:
k=0
# start
# 解码器中的 conv_block 和 ConvTranspose2dBlock 都使用了BN和relu虽然是不同人实现的，不过方法是一致的
layers += [ConvTranspose2dBlock(1037, 512, (4, 4), stride=2, padding=1, norm_fn=dec_norm_fn, acti_fn=dec_acti_fn)] # 0
layers += [Attention_block(F_g=512,F_l=512,F_int=256)] # 1 *
layers += [conv_block(ch_in=1024,ch_out=512)] # 2
layers += [ConvTranspose2dBlock(512+k, 256, (4, 4), stride=2, padding=1, norm_fn=dec_norm_fn, acti_fn=dec_acti_fn)] # 3
layers += [Attention_block(F_g=256,F_l=256,F_int=128)] # 4 *
layers += [conv_block(ch_in=512,ch_out=256)] # 5
layers += [ConvTranspose2dBlock(256+j, 128, (4, 4), stride=2, padding=1, norm_fn=dec_norm_fn, acti_fn=dec_acti_fn)] # 6
layers += [Self_Attn(128, 'relu')] # 7 ~
layers += [Attention_block(F_g=128,F_l=128,F_int=64)] # 8 *
layers += [conv_block(ch_in=256,ch_out=128)] # 9
layers += [ConvTranspose2dBlock(128, 64, (4, 4), stride=2, padding=1, norm_fn=dec_norm_fn, acti_fn=dec_acti_fn)] # 10
layers += [Self_Attn(64, 'relu')] # 11 ~
layers += [Attention_block(F_g=64,F_l=64,F_int=32)] # 12 *
layers += [conv_block(ch_in=128,ch_out=64)] # 13
layers += [ConvTranspose2dBlock(64, 3, (4, 4), stride=2, padding=1, norm_fn='none', acti_fn='tanh')] # 14
self.dec_layers = nn.ModuleList(layers)
def __init__(self,
dim=64,
norm_fn='instancenorm',
acti_fn='lrelu',
fc_dim=1024,
fc_norm_fn='none',
fc_acti_fn='lrelu',
n_layers=5,
img_size=128):
super(Discriminators, self).__init__()
self.f_size = img_size // 2**n_layers
layers = []
n_in = 3
for i in range(n_layers):
n_out = min(dim * 2**i, MAX_DIM)
layers += [
Conv2dBlock(n_in,
n_out, (4, 4),
stride=2,
padding=1,
norm_fn=norm_fn,
acti_fn=acti_fn)
]
n_in = n_out
self.conv = nn.Sequential(*layers)
self.fc_adv = nn.Sequential(
LinearBlock(1024 * self.f_size * self.f_size, fc_dim, fc_norm_fn,
fc_acti_fn), LinearBlock(fc_dim, 1, 'none', 'none'))
self.fc_cls = nn.Sequential(
LinearBlock(1024 * self.f_size * self.f_size, fc_dim, fc_norm_fn,
fc_acti_fn), LinearBlock(fc_dim, 13, 'none', 'none'))
def __init__(self, enc_dim=64, enc_layers=5, enc_norm_fn='batchnorm', enc_acti_fn='lrelu',
dec_dim=64, dec_layers=5, dec_norm_fn='batchnorm', dec_acti_fn='relu',
n_attrs=13, shortcut_layers=1, inject_layers=0, img_size=128):
super(Generator, self).__init__()
self.shortcut_layers = min(shortcut_layers, dec_layers - 1)
self.inject_layers = min(inject_layers, dec_layers - 1)
self.f_size = img_size // 2**enc_layers # f_size = 4 for 128x128
layers = []
n_in = 3
for i in range(enc_layers):
n_out = min(enc_dim * 2**i, MAX_DIM)
layers += [Conv2dBlock(
n_in, n_out, (4, 4), stride=2, padding=1, norm_fn=enc_norm_fn, acti_fn=enc_acti_fn
)]
n_in = n_out
self.enc_layers = nn.ModuleList(layers)
layers = []
n_in = n_in + n_attrs # 1024 + 13
for i in range(dec_layers):
if i < dec_layers - 1:
n_out = min(dec_dim * 2**(dec_layers-i-1), MAX_DIM)
layers += [ConvTranspose2dBlock(
n_in, n_out, (4, 4), stride=2, padding=1, norm_fn=dec_norm_fn, acti_fn=dec_acti_fn
)]
n_in = n_out
n_in = n_in + n_in//2 if self.shortcut_layers > i else n_in
n_in = n_in + n_attrs if self.inject_layers > i else n_in
else:
layers += [ConvTranspose2dBlock(
n_in, 3, (4, 4), stride=2, padding=1, norm_fn='none', acti_fn='tanh'
)]
self.dec_layers = nn.ModuleList(layers)
def __init__(self, enc_dim=64, enc_layers=5, enc_norm_fn='batchnorm', enc_acti_fn='lrelu',
dec_dim=64, dec_layers=5, dec_norm_fn='batchnorm', dec_acti_fn='relu',
n_attrs=13, shortcut_layers=1, inject_layers=0, img_size=128, use_stu=True):
super(Generator, self).__init__()
self.shortcut_layers = min(shortcut_layers, dec_layers - 1)
self.inject_layers = min(inject_layers, dec_layers - 1)
self.f_size = img_size // 2**enc_layers # f_size = 4 for 128x128
self.n_attrs = n_attrs
self.use_stu = use_stu
self.dec_layers = dec_layers
layers = []
n_in = 3
for i in range(enc_layers):
n_out = min(enc_dim * 2**i, MAX_DIM)
layers += [Conv2dBlock(
n_in, n_out, (4, 4), stride=2, padding=1, norm_fn=enc_norm_fn, acti_fn=enc_acti_fn
)]
n_in = n_out
self.Encode = nn.ModuleList(layers)
if use_stu:
layers = []
for i in reversed(range(enc_layers-1-self.shortcut_layers, enc_layers-1)):
conv_dim = min(enc_dim * 2**i, MAX_DIM)
layers += [ConvGRUCell(n_attrs, conv_dim, conv_dim, 3)]
self.stu = nn.ModuleList(layers)
layers = []
n_in = n_in + n_attrs # 1024 + 13
for i in range(dec_layers):
if i < dec_layers - 1:
n_out = min(dec_dim * 2 ** (dec_layers - i - 1), MAX_DIM)
if i == 0:
layers += [ConvTranspose2dBlock(
n_in, n_out, (4, 4), stride=2, padding=1, norm_fn=dec_norm_fn, acti_fn=dec_acti_fn
)]
elif i <= self.shortcut_layers:
n_in = n_in + n_in//2 if self.shortcut_layers >= i else n_in
layers += [ConvTranspose2dBlock(
n_in, n_out, (4, 4), stride=2, padding=1, norm_fn=dec_norm_fn, acti_fn=dec_acti_fn
)]
else:
layers += [ConvTranspose2dBlock(
n_in, n_out, (4, 4), stride=2, padding=1, norm_fn=dec_norm_fn, acti_fn=dec_acti_fn
)]
n_in = n_out
else:
layers += [ConvTranspose2dBlock(
n_in, n_in // 4, (4, 4), stride=2, padding=1, norm_fn=dec_norm_fn, acti_fn=dec_acti_fn
)]
layers += [ConvTranspose2dBlock(
n_in // 4, 3, (3, 3), stride=1, padding=1, norm_fn='none', acti_fn='tanh'
)]
self.Decode = nn.ModuleList(layers)
def __init__(self, dim=64, norm_fn='instancenorm', acti_fn='lrelu',
fc_dim=1024, fc_norm_fn='none', fc_acti_fn='lrelu', n_layers=2):
super(Discriminators, self).__init__()
layers = []
n_in = 256 * BasicBlock.expansion
for i in range(3, 3 + n_layers):
n_out = min(dim * 2**i, MAX_DIM)
layers += [Conv2dBlock(
n_in, n_out, 4, stride=2, padding=1, norm_fn=norm_fn, acti_fn=acti_fn
)]
n_in = n_out
self.conv = nn.Sequential(*layers)
self.fc_adv = nn.Sequential(
LinearBlock(4 * fc_dim, fc_dim, fc_norm_fn, fc_acti_fn),
LinearBlock(fc_dim, 1, 'none', 'none')
)
self.att_cls = ResNet(block=BasicBlock, layers=[2, 2, 2, 2], num_classes=13)
def __init__(self,
enc_dim=64,
enc_layers=5,
enc_norm_fn='batchnorm',
enc_acti_fn='lrelu',
dec_dim=64,
dec_layers=5,
dec_norm_fn='batchnorm',
dec_acti_fn='relu',
n_attrs=13,
shortcut_layers=0,
inject_layers=0,
img_size=128):
super(Generator, self).__init__()
self.shortcut_layers = min(shortcut_layers, dec_layers - 1)
self.inject_layers = min(inject_layers, dec_layers - 1)
self.f_size = img_size // 2**enc_layers # f_size = 4 for 128x128
layers = []
n_in = 3
# start
layers += [
Conv2dBlock(3,
64, (4, 4),
stride=2,
padding=1,
norm_fn=enc_norm_fn,
acti_fn=enc_acti_fn)
] # 0
layers += [Self_Attn(64, 'relu')] # 1
layers += [
Conv2dBlock(64,
128, (4, 4),
stride=2,
padding=1,
norm_fn=enc_norm_fn,
acti_fn=enc_acti_fn)
] # 2
layers += [Self_Attn(128, 'relu')] # 3
layers += [
Conv2dBlock(128,
256, (4, 4),
stride=2,
padding=1,
norm_fn=enc_norm_fn,
acti_fn=enc_acti_fn)
] # 4
layers += [
Conv2dBlock(256,
512, (4, 4),
stride=2,
padding=1,
norm_fn=enc_norm_fn,
acti_fn=enc_acti_fn)
] # 5
layers += [
Conv2dBlock(512,
1024, (4, 4),
stride=2,
padding=1,
norm_fn=enc_norm_fn,
acti_fn=enc_acti_fn)
] # 6
self.enc_layers = nn.ModuleList(layers)
# end
layers = []
n_in = n_in + n_attrs # 1024 + 13
if self.inject_layers >= 1:
k = 13
else:
k = 0
# start
layers += [
ConvTranspose2dBlock(1037,
1024, (4, 4),
stride=2,
padding=1,
norm_fn=dec_norm_fn,
acti_fn=dec_acti_fn)
] # 0
layers += [Attention_block(F_g=1024, F_l=512, F_int=512)] # 1
layers += [conv_block(ch_in=1536, ch_out=1024)] # 2
layers += [
ConvTranspose2dBlock(1024 + k,
512, (4, 4),
stride=2,
padding=1,
norm_fn=dec_norm_fn,
acti_fn=dec_acti_fn)
] # 3
layers += [Attention_block(F_g=512, F_l=256, F_int=256)] # 4
layers += [conv_block(ch_in=768, ch_out=512)] # 5
layers += [
ConvTranspose2dBlock(512,
256, (4, 4),
stride=2,
padding=1,
norm_fn=dec_norm_fn,
acti_fn=dec_acti_fn)
] # 6
layers += [Self_Attn(256, 'relu')] # 7
layers += [
ConvTranspose2dBlock(256,
128, (4, 4),
stride=2,
padding=1,
norm_fn=dec_norm_fn,
acti_fn=dec_acti_fn)
] # 8
layers += [Self_Attn(128, 'relu')] # 9
layers += [
ConvTranspose2dBlock(128,
3, (4, 4),
stride=2,
padding=1,
norm_fn='none',
acti_fn='tanh')
] # 10
self.dec_layers = nn.ModuleList(layers)
def __init__(self,
enc_dim=64,
enc_layers=5,
enc_norm_fn='batchnorm',
enc_acti_fn='lrelu',
dec_dim=64,
dec_layers=5,
dec_norm_fn='batchnorm',
dec_acti_fn='relu',
n_attrs=13,
shortcut_layers=0,
inject_layers=0,
img_size=128):
super(Generator, self).__init__()
self.shortcut_layers = min(shortcut_layers, dec_layers - 1)
self.inject_layers = min(inject_layers, dec_layers - 1)
self.f_size = img_size // 2**enc_layers # f_size = 4 for 128x128
layers = []
n_in = 3
for i in range(enc_layers):
n_out = min(enc_dim * 2**i, MAX_DIM)
layers += [
Conv2dBlock(n_in,
n_out, (4, 4),
stride=2,
padding=1,
norm_fn=enc_norm_fn,
acti_fn=enc_acti_fn)
]
n_in = n_out
self.enc_layers = nn.ModuleList(layers)
layers = []
n_in = n_in + n_attrs # 1024 + 13
for i in range(dec_layers):
if i < dec_layers - 3: # 0,1
n_out = min(dec_dim * 2**(dec_layers - i - 1), MAX_DIM)
layers += [
ConvTranspose2dBlock(n_in,
n_out, (4, 4),
stride=2,
padding=1,
norm_fn=dec_norm_fn,
acti_fn=dec_acti_fn)
]
layers += [
Attention_block(F_g=n_out,
F_l=n_out // 2,
F_int=n_out // 2)
] # attention模块输入F_g为输入的一半
layers += [conv_block(ch_in=int(1.5 * n_out), ch_out=n_out)]
n_in = n_out
n_in = n_in + n_attrs if self.inject_layers > i else n_in
elif i >= dec_layers - 3 and i < dec_layers - 1:
n_out = min(dec_dim * 2**(dec_layers - i - 1), MAX_DIM)
layers += [
ConvTranspose2dBlock(n_in,
n_out, (4, 4),
stride=2,
padding=1,
norm_fn=dec_norm_fn,
acti_fn=dec_acti_fn)
]
n_in = n_out
elif i == dec_layers - 1:
layers += [
ConvTranspose2dBlock(n_in,
3, (4, 4),
stride=2,
padding=1,
norm_fn='none',
acti_fn='tanh')
]
self.dec_layers = nn.ModuleList(layers)
