def __init__(self, opt):
super(SPADEModel, self).__init__(opt)
self.model_names = ['G']
self.visual_names = ['labels', 'fake_B', 'real_B']
self.modules = SPADEModelModules(opt).to(self.device)
if len(opt.gpu_ids) > 0:
self.modules = DataParallelWithCallback(self.modules, device_ids=opt.gpu_ids)
self.modules_on_one_gpu = self.modules.module
else:
self.modules_on_one_gpu = self.modules
if opt.isTrain:
self.model_names.append('D')
self.loss_names = ['G_gan', 'G_feat', 'G_vgg', 'D_real', 'D_fake']
self.optimizer_G, self.optimizer_D = self.modules_on_one_gpu.create_optimizers()
self.optimizers = [self.optimizer_G, self.optimizer_D]
if not opt.no_fid:
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 and not opt.no_mIoU:
self.drn_model = DRNSeg('drn_d_105', 19, pretrained=False)
util.load_network(self.drn_model, opt.drn_path, verbose=False)
self.drn_model.to(self.device)
self.drn_model.eval()
self.eval_dataloader = create_eval_dataloader(self.opt)
self.best_fid = 1e9
self.best_mIoU = -1e9
self.fids, self.mIoUs = [], []
self.is_best = False
self.npz = np.load(opt.real_stat_path)
else:
self.modules.eval()
self.train_dataloader = create_train_dataloader(opt)
def __init__(self, opt):
super(SPADEModel, self).__init__(opt)
self.model_names = ['G_student', 'G_teacher', 'D']
self.visual_names = ['labels', 'Tfake_B', 'Sfake_B', 'real_B']
self.model_names.append('D')
self.loss_names = [
'G_gan', 'G_feat', 'G_vgg', 'G_distill', 'D_real', 'D_fake'
]
if hasattr(opt, 'distiller'):
self.modules = SPADEDistillerModules(opt).to(self.device)
if len(opt.gpu_ids) > 0:
self.modules = DataParallelWithCallback(self.modules,
device_ids=opt.gpu_ids)
self.modules_on_one_gpu = self.modules.module
else:
self.modules_on_one_gpu = self.modules
for i in range(len(self.modules_on_one_gpu.mapping_layers)):
self.loss_names.append('G_distill%d' % i)
self.optimizer_G, self.optimizer_D = self.modules_on_one_gpu.create_optimizers(
)
self.optimizers = [self.optimizer_G, self.optimizer_D]
if not opt.no_fid:
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 and not opt.no_mIoU:
self.drn_model = DRNSeg('drn_d_105', 19, pretrained=False)
util.load_network(self.drn_model, opt.drn_path, verbose=False)
self.drn_model.to(self.device)
self.drn_model.eval()
self.eval_dataloader = create_eval_dataloader(self.opt)
self.best_fid = 1e9
self.best_mIoU = -1e9
self.fids, self.mIoUs = [], []
self.is_best = False
self.npz = np.load(opt.real_stat_path)
model_profiling(self.modules_on_one_gpu.netG_teacher,
self.opt.data_height,
self.opt.data_width,
channel=self.opt.data_channel,
num_forwards=0,
verbose=False)
model_profiling(self.modules_on_one_gpu.netG_student,
self.opt.data_height,
self.opt.data_width,
channel=self.opt.data_channel,
num_forwards=0,
verbose=False)
print(
f'Teacher FLOPs: {self.modules_on_one_gpu.netG_teacher.n_macs}, Student FLOPs: {self.modules_on_one_gpu.netG_student.n_macs}.'
)
def __init__(self, opt):
assert opt.isTrain
valid_netGs = [
'spade', 'mobile_spade', 'super_mobile_spade', 'sub_mobile_spade'
]
assert opt.teacher_netG in valid_netGs and opt.student_netG in valid_netGs
super(SPADEModel, self).__init__(opt)
self.model_names = ['G_student', 'G_teacher', 'D']
self.visual_names = ['labels', 'Tfake_B', 'Sfake_B', 'real_B']
self.model_names.append('D')
self.loss_names = [
'G_gan', 'G_feat', 'G_vgg', 'G_distill', 'D_real', 'D_fake'
]
if hasattr(opt, 'distiller'):
self.modules = SPADEDistillerModules(opt).to(self.device)
if len(opt.gpu_ids) > 0:
self.modules = DataParallelWithCallback(self.modules,
device_ids=opt.gpu_ids)
self.modules_on_one_gpu = self.modules.module
else:
self.modules_on_one_gpu = self.modules
else:
self.modules = SPADESupernetModules(opt).to(self.device)
if len(opt.gpu_ids) > 0:
self.modules = DataParallelWithCallback(self.modules,
device_ids=opt.gpu_ids)
self.modules_on_one_gpu = self.modules.module
else:
self.modules_on_one_gpu = self.modules
for i in range(len(self.modules_on_one_gpu.mapping_layers)):
self.loss_names.append('G_distill%d' % i)
self.optimizer_G, self.optimizer_D = self.modules_on_one_gpu.create_optimizers(
)
self.optimizers = [self.optimizer_G, self.optimizer_D]
if not opt.no_fid:
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 and not opt.no_mIoU:
self.drn_model = DRNSeg('drn_d_105', 19, pretrained=False)
util.load_network(self.drn_model, opt.drn_path, verbose=False)
self.drn_model.to(self.device)
self.drn_model.eval()
self.eval_dataloader = create_eval_dataloader(self.opt)
self.best_fid = 1e9
self.best_mIoU = -1e9
self.fids, self.mIoUs = [], []
self.is_best = False
self.npz = np.load(opt.real_stat_path)