def __init__(self, opt):
super(SuperMobileSPADEGenerator, self).__init__()
self.opt = opt
nf = opt.ngf
self.sw, self.sh = self.compute_latent_vector_size(opt)
# downsampled segmentation map instead of random z
self.fc = SuperConv2d(self.opt.semantic_nc, 16 * nf, 3, padding=1)
self.head_0 = SuperMobileSPADEResnetBlock(16 * nf, 16 * nf, opt)
self.G_middle_0 = SuperMobileSPADEResnetBlock(16 * nf, 16 * nf, opt)
self.G_middle_1 = SuperMobileSPADEResnetBlock(16 * nf, 16 * nf, opt)
self.up_0 = SuperMobileSPADEResnetBlock(16 * nf, 8 * nf, opt)
self.up_1 = SuperMobileSPADEResnetBlock(8 * nf, 4 * nf, opt)
self.up_2 = SuperMobileSPADEResnetBlock(4 * nf, 2 * nf, opt)
self.up_3 = SuperMobileSPADEResnetBlock(2 * nf, 1 * nf, opt)
final_nc = nf
if opt.num_upsampling_layers == 'most':
self.up_4 = SuperMobileSPADEResnetBlock(1 * nf, nf // 2, opt)
final_nc = nf // 2
self.conv_img = SuperConv2d(final_nc, 3, 3, padding=1)
self.up = nn.Upsample(scale_factor=2)
def __init__(self, fin, fout, opt):
super(SuperMobileSPADEResnetBlock, self).__init__()
# Attributes
self.learned_shortcut = (fin != fout)
fmiddle = min(fin, fout)
# create conv layers
self.conv_0 = SuperConv2d(fin, fmiddle, kernel_size=3, padding=1)
self.conv_1 = SuperConv2d(fmiddle, fout, kernel_size=3, padding=1)
if self.learned_shortcut:
self.conv_s = SuperConv2d(fin, fout, kernel_size=1, bias=False)
# define normalization layers
spade_config_str = opt.norm_G
self.norm_0 = SuperMobileSPADE(spade_config_str,
fin,
opt.semantic_nc,
nhidden=opt.ngf * 2)
self.norm_1 = SuperMobileSPADE(spade_config_str,
fmiddle,
opt.semantic_nc,
nhidden=opt.ngf * 2)
if self.learned_shortcut:
self.norm_s = SuperMobileSPADE(spade_config_str,
fin,
opt.semantic_nc,
nhidden=opt.ngf * 2)
def __init__(self, config_text, norm_nc, label_nc, nhidden=128):
super(SuperMobileSPADE, self).__init__()
assert config_text.startswith('spade')
parsed = re.search(r'spade(\D+)(\d)x\d', config_text)
param_free_norm_type = str(parsed.group(1))
ks = int(parsed.group(2))
if param_free_norm_type == 'instance':
self.param_free_norm = nn.InstanceNorm2d(norm_nc, affine=False)
elif param_free_norm_type == 'syncbatch':
self.param_free_norm = SuperSynchronizedBatchNorm2d(norm_nc,
affine=False)
elif param_free_norm_type == 'batch':
self.param_free_norm = nn.BatchNorm2d(norm_nc, affine=False)
else:
raise ValueError(
'%s is not a recognized param-free norm type in SPADE' %
param_free_norm_type)
# The dimension of the intermediate embedding space. Yes, hardcoded.
pw = ks // 2
self.mlp_shared = nn.Sequential(
SuperConv2d(label_nc, nhidden, kernel_size=ks, padding=pw),
nn.ReLU())
self.mlp_gamma = SuperSeparableConv2d(nhidden,
norm_nc,
kernel_size=ks,
padding=pw)
self.mlp_beta = SuperSeparableConv2d(nhidden,
norm_nc,
kernel_size=ks,
padding=pw)
def __init__(self, opt):
assert opt.isTrain
opt = copy.deepcopy(opt)
if len(opt.gpu_ids) > 0:
opt.gpu_ids = opt.gpu_ids[:1]
self.gpu_ids = opt.gpu_ids
super(SPADEModelModules, self).__init__()
self.opt = opt
self.model_names = ['G_student', 'G_teacher', 'D']
teacher_opt = self.create_option('teacher')
self.netG_teacher = networks.define_G(opt.teacher_netG,
gpu_ids=self.gpu_ids,
opt=teacher_opt)
student_opt = self.create_option('student')
self.netG_student = networks.define_G(opt.student_netG,
init_type=opt.init_type,
init_gain=opt.init_gain,
gpu_ids=self.gpu_ids,
opt=student_opt)
if hasattr(opt, 'distiller'):
pretrained_opt = self.create_option('pretrained')
self.netG_pretrained = networks.define_G(opt.pretrained_netG,
gpu_ids=self.gpu_ids,
opt=pretrained_opt)
self.netD = networks.define_D(opt.netD,
init_type=opt.init_type,
init_gain=opt.init_gain,
gpu_ids=self.gpu_ids,
opt=opt)
self.mapping_layers = ['head_0', 'G_middle_1', 'up_1']
self.netAs = nn.ModuleList()
for i, mapping_layer in enumerate(self.mapping_layers):
if mapping_layer != 'up_1':
fs, ft = opt.student_ngf * 16, opt.teacher_ngf * 16
else:
fs, ft = opt.student_ngf * 4, opt.teacher_ngf * 4
if hasattr(opt, 'distiller'):
netA = nn.Conv2d(in_channels=fs,
out_channels=ft,
kernel_size=1)
else:
netA = SuperConv2d(in_channels=fs,
out_channels=ft,
kernel_size=1)
networks.init_net(netA, opt.init_type, opt.init_gain, self.gpu_ids)
self.netAs.append(netA)
self.criterionGAN = GANLoss(opt.gan_mode)
self.criterionFeat = nn.L1Loss()
self.criterionVGG = VGGLoss()
self.optimizers = []
self.netG_teacher.eval()
self.config = None
def __init__(self, n_downsample, input_dim, dim, style_dim, norm, activ,
pad_type):
super(SuperStyleEncoder, self).__init__()
self.n_downsample = n_downsample
self.style_dim = style_dim
self.model = []
self.model += [
SuperConv2dBlock(input_dim,
dim,
7,
1,
3,
norm=norm,
activation=activ,
pad_type=pad_type)
]
for i in range(2):
self.model += [
SuperConv2dBlock(dim,
2 * dim,
4,
2,
1,
norm=norm,
activation=activ,
pad_type=pad_type)
]
dim *= 2
for i in range(n_downsample - 2):
self.model += [
SuperConv2dBlock(dim,
dim,
4,
2,
1,
norm=norm,
activation=activ,
pad_type=pad_type)
]
self.model += [nn.AdaptiveAvgPool2d(1)] # global average pooling
self.model += [SuperConv2d(dim, style_dim, 1, 1, 0)]
self.model = nn.Sequential(*self.model)
self.output_dim = dim
def __init__(self, opt):
assert opt.isTrain
valid_netGs = [
'munit', 'super_munit', 'super_mobile_munit',
'super_mobile_munit2', 'super_mobile_munit3'
]
assert opt.teacher_netG in valid_netGs and opt.student_netG in valid_netGs
super(BaseMunitDistiller, self).__init__(opt)
self.loss_names = [
'G_gan', 'G_rec_x', 'G_rec_c', 'G_rec_s', 'D_fake', 'D_real'
]
if not opt.student_no_style_encoder:
self.loss_names.append('G_rec_s')
self.optimizers = []
self.image_paths = []
self.visual_names = ['real_A', 'Sfake_B', 'Tfake_B', 'real_B']
self.model_names = ['netG_student', 'netG_teacher', 'netD']
opt_teacher = self.create_option('teacher')
self.netG_teacher = networks.define_G(opt.teacher_netG,
init_type=opt.init_type,
init_gain=opt.init_gain,
gpu_ids=self.gpu_ids,
opt=opt_teacher)
opt_student = self.create_option('student')
self.netG_student = networks.define_G(opt.student_netG,
init_type=opt.init_type,
init_gain=opt.init_gain,
gpu_ids=self.gpu_ids,
opt=opt_student)
self.netD = networks.define_D(opt.netD,
input_nc=opt.output_nc,
init_type='normal',
init_gain=opt.init_gain,
gpu_ids=self.gpu_ids,
opt=opt)
if hasattr(opt, 'distiller'):
self.netA = nn.Conv2d(in_channels=4 * opt.student_ngf,
out_channels=4 * opt.teacher_ngf,
kernel_size=1).to(self.device)
else:
self.netA = SuperConv2d(in_channels=4 * opt.student_ngf,
out_channels=4 * opt.teacher_ngf,
kernel_size=1).to(self.device)
networks.init_net(self.netA)
self.netG_teacher.eval()
self.criterionGAN = GANLoss(opt.gan_mode).to(self.device)
self.criterionRec = torch.nn.L1Loss()
G_params = []
G_params.append(self.netG_student.parameters())
G_params.append(self.netA.parameters())
self.optimizer_G = torch.optim.Adam(itertools.chain(*G_params),
lr=opt.lr,
betas=(opt.beta1, 0.999),
weight_decay=opt.weight_decay)
self.optimizer_D = torch.optim.Adam(self.netD.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999),
weight_decay=opt.weight_decay)
self.optimizers.append(self.optimizer_G)
self.optimizers.append(self.optimizer_D)
self.eval_dataloader = create_eval_dataloader(self.opt,
direction=opt.direction)
self.inception_model, _, _ = create_metric_models(opt,
device=self.device)
self.npz = np.load(opt.real_stat_path)
self.is_best = False
def __init__(self, opt):
assert opt.isTrain
super(BaseResnetDistiller, self).__init__(opt)
self.loss_names = ['G_gan', 'G_distill', 'G_recon', 'D_fake', 'D_real']
self.optimizers = []
self.image_paths = []
self.visual_names = ['real_A', 'Sfake_B', 'Tfake_B', 'real_B']
self.model_names = ['netG_student', 'netG_teacher', 'netD']
self.netG_teacher = networks.define_G(opt.input_nc, opt.output_nc, opt.teacher_ngf,
opt.teacher_netG, opt.norm, opt.teacher_dropout_rate,
opt.init_type, opt.init_gain, self.gpu_ids, opt=opt)
self.netG_student = networks.define_G(opt.input_nc, opt.output_nc, opt.student_ngf,
opt.student_netG, opt.norm, opt.student_dropout_rate,
opt.init_type, opt.init_gain, self.gpu_ids, opt=opt)
if getattr(opt, 'sort_channels', False) and opt.restore_student_G_path is not None:
self.netG_student_tmp = networks.define_G(opt.input_nc, opt.output_nc, opt.student_ngf,
opt.student_netG.replace('super_', ''), opt.norm,
opt.student_dropout_rate, opt.init_type, opt.init_gain,
self.gpu_ids, opt=opt)
if hasattr(opt, 'distiller'):
self.netG_pretrained = networks.define_G(opt.input_nc, opt.output_nc, opt.pretrained_ngf,
opt.pretrained_netG, opt.norm, 0,
opt.init_type, opt.init_gain, self.gpu_ids, opt=opt)
if opt.dataset_mode == 'aligned':
self.netD = networks.define_D(opt.input_nc + opt.output_nc, opt.ndf, opt.netD,
opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids)
elif opt.dataset_mode == 'unaligned':
self.netD = networks.define_D(opt.output_nc, opt.ndf, opt.netD,
opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids)
else:
raise NotImplementedError('Unknown dataset mode [%s]!!!' % opt.dataset_mode)
self.netG_teacher.eval()
self.criterionGAN = models.modules.loss.GANLoss(opt.gan_mode).to(self.device)
if opt.recon_loss_type == 'l1':
self.criterionRecon = torch.nn.L1Loss()
elif opt.recon_loss_type == 'l2':
self.criterionRecon = torch.nn.MSELoss()
elif opt.recon_loss_type == 'smooth_l1':
self.criterionRecon = torch.nn.SmoothL1Loss()
elif opt.recon_loss_type == 'vgg':
self.criterionRecon = models.modules.loss.VGGLoss(self.device)
else:
raise NotImplementedError('Unknown reconstruction loss type [%s]!' % opt.loss_type)
if isinstance(self.netG_teacher, nn.DataParallel):
self.mapping_layers = ['module.model.%d' % i for i in range(9, 21, 3)]
else:
self.mapping_layers = ['model.%d' % i for i in range(9, 21, 3)]
self.netAs = []
self.Tacts, self.Sacts = {}, {}
G_params = [self.netG_student.parameters()]
for i, n in enumerate(self.mapping_layers):
ft, fs = self.opt.teacher_ngf, self.opt.student_ngf
if hasattr(opt, 'distiller'):
netA = nn.Conv2d(in_channels=fs * 4, out_channels=ft * 4, kernel_size=1). \
to(self.device)
else:
netA = SuperConv2d(in_channels=fs * 4, out_channels=ft * 4, kernel_size=1). \
to(self.device)
networks.init_net(netA)
G_params.append(netA.parameters())
self.netAs.append(netA)
self.loss_names.append('G_distill%d' % i)
self.optimizer_G = torch.optim.Adam(itertools.chain(*G_params), lr=opt.lr, betas=(opt.beta1, 0.999))
self.optimizer_D = torch.optim.Adam(self.netD.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
self.optimizers.append(self.optimizer_G)
self.optimizers.append(self.optimizer_D)
self.eval_dataloader = create_eval_dataloader(self.opt, direction=opt.direction)
block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[2048]
self.inception_model = InceptionV3([block_idx])
self.inception_model.to(self.device)
self.inception_model.eval()
if 'cityscapes' in opt.dataroot:
self.drn_model = DRNSeg('drn_d_105', 19, pretrained=False)
util.load_network(self.drn_model, opt.drn_path, verbose=False)
if len(opt.gpu_ids) > 0:
self.drn_model.to(self.device)
self.drn_model = nn.DataParallel(self.drn_model, opt.gpu_ids)
self.drn_model.eval()
self.npz = np.load(opt.real_stat_path)
self.is_best = False
def __init__(self,
input_nc,
output_nc,
ngf,
norm_layer=nn.BatchNorm2d,
dropout_rate=0,
n_blocks=6,
padding_type='reflect'):
assert n_blocks >= 0
super(SuperMobileResnetGenerator, self).__init__()
if type(norm_layer) == functools.partial:
use_bias = norm_layer.func == nn.InstanceNorm2d
else:
use_bias = norm_layer == nn.InstanceNorm2d
model = [
nn.ReflectionPad2d(3),
SuperConv2d(input_nc, ngf, kernel_size=7, padding=0,
bias=use_bias),
norm_layer(ngf),
nn.ReLU(True)
]
n_downsampling = 2
for i in range(n_downsampling): # add downsampling layers
mult = 2**i
model += [
SuperConv2d(ngf * mult,
ngf * mult * 2,
kernel_size=3,
stride=2,
padding=1,
bias=use_bias),
norm_layer(ngf * mult * 2),
nn.ReLU(True)
]
mult = 2**n_downsampling
n_blocks1 = n_blocks // 3
n_blocks2 = n_blocks1
n_blocks3 = n_blocks - n_blocks1 - n_blocks2
for i in range(n_blocks1):
model += [
SuperMobileResnetBlock(ngf * mult,
padding_type=padding_type,
norm_layer=norm_layer,
dropout_rate=dropout_rate,
use_bias=use_bias)
]
for i in range(n_blocks2):
model += [
SuperMobileResnetBlock(ngf * mult,
padding_type=padding_type,
norm_layer=norm_layer,
dropout_rate=dropout_rate,
use_bias=use_bias)
]
for i in range(n_blocks3):
model += [
SuperMobileResnetBlock(ngf * mult,
padding_type=padding_type,
norm_layer=norm_layer,
dropout_rate=dropout_rate,
use_bias=use_bias)
]
for i in range(n_downsampling): # add upsampling layers
mult = 2**(n_downsampling - i)
model += [
SuperConvTranspose2d(ngf * mult,
int(ngf * mult / 2),
kernel_size=3,
stride=2,
padding=1,
output_padding=1,
bias=use_bias),
norm_layer(int(ngf * mult / 2)),
nn.ReLU(True)
]
model += [nn.ReflectionPad2d(3)]
model += [SuperConv2d(ngf, output_nc, kernel_size=7, padding=0)]
model += [nn.Tanh()]
self.model = nn.Sequential(*model)
def __init__(self,
input_dim,
output_dim,
kernel_size,
stride,
padding=0,
norm='none',
activation='relu',
pad_type='zero'):
super(SuperConv2dBlock, self).__init__()
self.use_bias = True
# initialize padding
if pad_type == 'reflect':
self.pad = nn.ReflectionPad2d(padding)
elif pad_type == 'replicate':
self.pad = nn.ReplicationPad2d(padding)
elif pad_type == 'zero':
self.pad = nn.ZeroPad2d(padding)
else:
assert 0, "Unsupported padding type: {}".format(pad_type)
# initialize normalization
norm_dim = output_dim
if norm == 'bn':
raise NotImplementedError
elif norm == 'in':
self.norm = nn.InstanceNorm2d(norm_dim)
elif norm == 'ln':
self.norm = SuperLayerNorm(norm_dim)
elif norm == 'adain':
self.norm = SuperAdaptiveInstanceNorm2d(norm_dim)
elif norm == 'none' or norm == 'sn':
self.norm = None
else:
assert 0, "Unsupported normalization: {}".format(norm)
# initialize activation
if activation == 'relu':
self.activation = nn.ReLU(inplace=True)
elif activation == 'lrelu':
self.activation = nn.LeakyReLU(0.2, inplace=True)
elif activation == 'prelu':
self.activation = nn.PReLU()
elif activation == 'selu':
self.activation = nn.SELU(inplace=True)
elif activation == 'tanh':
self.activation = nn.Tanh()
elif activation == 'none':
self.activation = None
else:
assert 0, "Unsupported activation: {}".format(activation)
# initialize convolution
if norm == 'sn':
raise NotImplementedError
else:
self.conv = SuperConv2d(input_dim,
output_dim,
kernel_size,
stride,
bias=self.use_bias)
def __init__(self, opt):
assert opt.isTrain
valid_netGs = [
'resnet_9blocks', 'mobile_resnet_9blocks',
'super_mobile_resnet_9blocks', 'sub_mobile_resnet_9blocks'
]
assert opt.teacher_netG in valid_netGs and opt.student_netG in valid_netGs
super(BaseResnetDistiller, self).__init__(opt)
self.loss_names = ['G_gan', 'G_distill', 'G_recon', 'D_fake', 'D_real']
self.optimizers = []
self.image_paths = []
self.visual_names = ['real_A', 'Sfake_B', 'Tfake_B', 'real_B']
self.model_names = ['netG_student', 'netG_teacher', 'netD']
self.netG_teacher = networks.define_G(
opt.teacher_netG,
input_nc=opt.input_nc,
output_nc=opt.output_nc,
ngf=opt.teacher_ngf,
norm=opt.norm,
dropout_rate=opt.teacher_dropout_rate,
gpu_ids=self.gpu_ids,
opt=opt)
self.netG_student = networks.define_G(
opt.student_netG,
input_nc=opt.input_nc,
output_nc=opt.output_nc,
ngf=opt.student_ngf,
norm=opt.norm,
dropout_rate=opt.student_dropout_rate,
init_type=opt.init_type,
init_gain=opt.init_gain,
gpu_ids=self.gpu_ids,
opt=opt)
if hasattr(opt, 'distiller'):
self.netG_pretrained = networks.define_G(opt.pretrained_netG,
input_nc=opt.input_nc,
output_nc=opt.output_nc,
ngf=opt.pretrained_ngf,
norm=opt.norm,
gpu_ids=self.gpu_ids,
opt=opt)
if opt.dataset_mode == 'aligned':
self.netD = networks.define_D(opt.netD,
input_nc=opt.input_nc +
opt.output_nc,
ndf=opt.ndf,
n_layers_D=opt.n_layers_D,
norm=opt.norm,
init_type=opt.init_type,
init_gain=opt.init_gain,
gpu_ids=self.gpu_ids,
opt=opt)
elif opt.dataset_mode == 'unaligned':
self.netD = networks.define_D(opt.netD,
input_nc=opt.output_nc,
ndf=opt.ndf,
n_layers_D=opt.n_layers_D,
norm=opt.norm,
init_type=opt.init_type,
init_gain=opt.init_gain,
gpu_ids=self.gpu_ids,
opt=opt)
else:
raise NotImplementedError('Unknown dataset mode [%s]!!!' %
opt.dataset_mode)
self.netG_teacher.eval()
self.criterionGAN = GANLoss(opt.gan_mode).to(self.device)
if opt.recon_loss_type == 'l1':
self.criterionRecon = torch.nn.L1Loss()
elif opt.recon_loss_type == 'l2':
self.criterionRecon = torch.nn.MSELoss()
elif opt.recon_loss_type == 'smooth_l1':
self.criterionRecon = torch.nn.SmoothL1Loss()
elif opt.recon_loss_type == 'vgg':
self.criterionRecon = models.modules.loss.VGGLoss(self.device)
else:
raise NotImplementedError(
'Unknown reconstruction loss type [%s]!' % opt.loss_type)
if isinstance(self.netG_teacher, nn.DataParallel):
self.mapping_layers = [
'module.model.%d' % i for i in range(9, 21, 3)
]
else:
self.mapping_layers = ['model.%d' % i for i in range(9, 21, 3)]
self.netAs = []
self.Tacts, self.Sacts = {}, {}
G_params = [self.netG_student.parameters()]
for i, n in enumerate(self.mapping_layers):
ft, fs = self.opt.teacher_ngf, self.opt.student_ngf
if hasattr(opt, 'distiller'):
netA = nn.Conv2d(in_channels=fs * 4, out_channels=ft * 4, kernel_size=1). \
to(self.device)
else:
netA = SuperConv2d(in_channels=fs * 4, out_channels=ft * 4, kernel_size=1). \
to(self.device)
networks.init_net(netA)
G_params.append(netA.parameters())
self.netAs.append(netA)
self.loss_names.append('G_distill%d' % i)
self.optimizer_G = torch.optim.Adam(itertools.chain(*G_params),
lr=opt.lr,
betas=(opt.beta1, 0.999))
self.optimizer_D = torch.optim.Adam(self.netD.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
self.optimizers.append(self.optimizer_G)
self.optimizers.append(self.optimizer_D)
self.eval_dataloader = create_eval_dataloader(self.opt,
direction=opt.direction)
self.inception_model, self.drn_model, _ = create_metric_models(
opt, device=self.device)
self.npz = np.load(opt.real_stat_path)
self.is_best = False