def test_cyclegan_fid(model, opt): opt.phase = 'test' opt.num_threads = 0 opt.batch_size = 1 opt.serial_batches = True opt.no_flip = True opt.load_size = 256 opt.display_id = -1 dataset = create_dataset(opt) model.model_eval() result_dir = os.path.join(opt.checkpoints_dir, opt.name, 'test_results') util.mkdirs(result_dir) fake_A = {} fake_B = {} for i, data in enumerate(dataset): model.set_input(data) with torch.no_grad(): model.forward() visuals = model.get_current_visuals() fake_B[data['A_paths'][0]] = visuals['fake_B'] fake_A[data['B_paths'][0]] = visuals['fake_A'] util.save_images(visuals, model.image_paths, result_dir, direction=opt.direction, aspect_ratio=opt.aspect_ratio) # print('Calculating AtoB FID...', flush=True) block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[2048] inception_model = InceptionV3([block_idx]) inception_model.to(model.device) inception_model.eval() npz = np.load(os.path.join(opt.dataroot, 'real_stat_B.npz')) AtoB_fid = get_fid(list(fake_B.values()), inception_model, npz, model.device, opt.batch_size) block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[2048] inception_model = InceptionV3([block_idx]) inception_model.to(model.device) inception_model.eval() npz = np.load(os.path.join(opt.dataroot, 'real_stat_A.npz')) BtoA_fid = get_fid(list(fake_A.values()), inception_model, npz, model.device, opt.batch_size) return AtoB_fid, BtoA_fid
def main(opt): dataloader = create_dataloader(opt) device = torch.device('cuda:{}'.format(opt.gpu_ids[0])) if opt.gpu_ids \ else torch.device('cpu') block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[2048] inception_model = InceptionV3([block_idx]) inception_model.to(device) inception_model.eval() tensors = [] for i, data_i in enumerate(tqdm.tqdm(dataloader)): if opt.dataset_mode in ['single', 'aligned']: tensor = data_i[opt.direction[-1]] else: tensor = data_i['image'] tensors.append(tensor) tensors = torch.cat(tensors, dim=0) tensors = util.tensor2im(tensors).astype(float) mu, sigma = _compute_statistics_of_ims(tensors, inception_model, 32, 2048, device, use_tqdm=True) np.savez(opt.output_path, mu=mu, sigma=sigma)
def create_metric_models(opt, device): if not opt.no_fid: block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[2048] inception_model = InceptionV3([block_idx]) if len(opt.gpu_ids) > 1: inception_model = nn.DataParallel(inception_model, opt.gpu_ids) inception_model.to(device) inception_model.eval() else: inception_model = None if 'cityscapes' in opt.dataroot and opt.direction == 'BtoA': drn_model = DRNSeg('drn_d_105', 19, pretrained=False) util.load_network(drn_model, opt.drn_path, verbose=False) if len(opt.gpu_ids) > 0: drn_model = nn.DataParallel(drn_model, opt.gpu_ids) drn_model.to(device) drn_model.eval() else: drn_model = None if 'coco' in opt.dataroot and not opt.no_mIoU and opt.direction == 'BtoA': deeplabv2_model = MSC(DeepLabV2(n_classes=182, n_blocks=[3, 4, 23, 3], atrous_rates=[6, 12, 18, 24]), scales=[0.5, 0.75]) util.load_network(deeplabv2_model, opt.deeplabv2_path, verbose=False) if len(opt.gpu_ids) > 1: deeplabv2_model = nn.DataParallel(deeplabv2_model, opt.gpu_ids) deeplabv2_model.to(device) deeplabv2_model.eval() else: deeplabv2_model = None return inception_model, drn_model, deeplabv2_model
def main(cfgs): fluid.enable_imperative() if 'resnet' in cfgs.netG: from configs.resnet_configs import get_configs else: raise NotImplementedError configs = get_configs(config_name=cfgs.config_set) configs = list(configs.all_configs()) data_loader, id2name = create_eval_data(cfgs, direction=cfgs.direction) model = TestModel(cfgs) model.setup() ### load_network ### this input used in compute model flops and params for data in data_loader: model.set_input(data) break npz = np.load(cfgs.real_stat_path) results = [] for config in configs: fakes, names = [], [] flops, _ = model.profile(config=config) s_time = time.time() for i, data in enumerate(data_loader()): model.set_input(data) model.test(config) generated = model.fake_B fakes.append(generated.detach().numpy()) name = id2name[i] save_path = os.path.join(cfgs.save_dir, 'test' + str(config)) if not os.path.exists(save_path): os.makedirs(save_path) save_path = os.path.join(save_path, name) names.append(name) if i < cfgs.num_test: image = util.tensor2img(generated) util.save_image(image, save_path) result = {'config_str': encode_config(config), 'flops': flops} ### compute FLOPs fluid.disable_imperative() if not cfgs.no_fid: block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[2048] inception_model = InceptionV3([block_idx]) fid = get_fid(fakes, inception_model, npz, cfgs.inception_model_path, batch_size=cfgs.batch_size, use_gpu=cfgs.use_gpu) result['fid'] = fid fluid.enable_imperative() e_time = (time.time() - s_time) / 60 result['time'] = e_time print(result) results.append(result) if not os.path.exists(cfgs.save_dir): os.makedirs(os.path.dirname(cfgs.save_dir)) save_file = os.path.join(cfgs.save_dir, 'search_result.pkl') with open(save_file, 'wb') as f: pickle.dump(results, f) print('Successfully finish searching!!!')
def main(configs, opt, gpu_id, queue, verbose): opt.gpu_ids = [gpu_id] dataloader = create_dataloader(opt, verbose) model = create_model(opt, verbose) model.setup(opt, verbose) device = model.device if not opt.no_fid: block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[2048] inception_model = InceptionV3([block_idx]) inception_model.to(device) inception_model.eval() if 'cityscapes' in opt.dataroot and opt.direction == 'BtoA': drn_model = DRNSeg('drn_d_105', 19, pretrained=False) util.load_network(drn_model, opt.drn_path, verbose=False) if len(opt.gpu_ids) > 0: drn_model = nn.DataParallel(drn_model, opt.gpu_ids) drn_model.eval() npz = np.load(opt.real_stat_path) results = [] for config in tqdm.tqdm(configs): fakes, names = [], [] for i, data_i in enumerate(dataloader): model.set_input(data_i) if i == 0: macs, _ = model.profile(config) model.test(config) fakes.append(model.fake_B.cpu()) for path in model.get_image_paths(): short_path = ntpath.basename(path) name = os.path.splitext(short_path)[0] names.append(name) result = {'config_str': encode_config(config), 'macs': macs} if not opt.no_fid: fid = get_fid(fakes, inception_model, npz, device, opt.batch_size, use_tqdm=False) result['fid'] = fid if 'cityscapes' in opt.dataroot and opt.direction == 'BtoA': mAP = get_mAP(fakes, names, drn_model, device, data_dir=opt.cityscapes_path, batch_size=opt.batch_size, num_workers=opt.num_threads, use_tqdm=False) result['mAP'] = mAP print(result, flush=True) # print('Time Cost: %.2fmin' % ((time.time() - start_time) / 60), flush=True) results.append(result) queue.put(results)
def main(cfgs): images = read_data(cfgs.dataroot, cfgs.filename) block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[2048] inception_model = InceptionV3([block_idx]) images = np.concatenate(images, axis=0) images = util.tensor2img(images).astype('float32') m2, s2 = _compute_statistic_of_img(images, inception_model, 32, 2048, cfgs.use_gpu, cfgs.inception_model_path) np.savez(cfgs.save_dir, mu=m2, sigma=s2)
def main(img_folder_path): inception_model = InceptionV3() # inception_model.cuda() m_r_A, C_r_A = mean_var_cal(img_folder_path, 'real_A', inception_model) m_f_A, C_f_A = mean_var_cal(img_folder_path, 'fake_A', inception_model) m_r_B, C_r_B = mean_var_cal(img_folder_path, 'real_B', inception_model) m_f_B, C_f_B = mean_var_cal(img_folder_path, 'fake_B', inception_model) fid_value_A = fid(m_r_A, m_f_A, C_r_A, C_f_A) fid_value_B = fid(m_r_B, m_f_B, C_r_B, C_f_B) print(args.name, 'fid_value_A', fid_value_A, 'fid_value_B', fid_value_B) return fid_value_A, fid_value_B
def get_FIDScore(images): model = InceptionV3() model.cuda() model.eval() num_image = len(images) fake = np.empty((num_image, 2048)) real = np.empty((num_image, 2048)) for im in range(num_image): fake[im] = model(images[im]['fake'])[0].reshape(1, -1).cpu().data.numpy() real[im] = model(images[im]['real'])[0].reshape(1, -1).cpu().data.numpy() #fake = model(image['fake']) #real = model(image['real']) #fake = fake[0].reshape(fake[0].shape[0], -1).cpu().data.numpy() #real = real[0].reshape(real[0].shape[0], -1).cpu().data.numpy() #u_fake = fake.mean() #u_real = real.mean() u_fake = np.mean(fake, axis=0) u_real = np.mean(real, axis=0) #np_fake = fake.cpu().data.numpy() #np_real = real.cpu().data.numpy() sig1 = np.cov(fake, rowvar=False) #sig1 = np.cov(np_fake) sig2 = np.cov(real, rowvar=False) #sig2 = np.cov(np_real) diff = u_fake - u_real # use sqrtm to do square root for matrix covmean, _ = linalg.sqrtm(sig1.dot(sig2), disp=False) #covmean = np.sqrt(sig1 * sig2) # if covmean has complex value, change to real number if np.iscomplexobj(covmean): covmean = covmean.real tr_cov = np.trace(covmean) #covmean = np.trace(covmean) #covmean = torch.from_numpy(covmean).type(diff.type).to(diff.device) #return diff.dot(diff) + np.trace(sig1) + np.trace(sig2) - 2 * covmean return diff.dot(diff) + np.trace(sig1) + np.trace(sig2) - 2 * tr_cov
def calculate_fid_given_ims(ims_fake, ims_real, batch_size, cuda, dims, model=None, use_tqdm=False): if model is None: block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[dims] model = InceptionV3([block_idx]) if cuda: model.cuda() m1, s1 = _compute_statistics_of_ims(ims_fake, model, batch_size, dims, cuda, use_tqdm=use_tqdm) m2, s2 = _compute_statistics_of_ims(ims_real, model, batch_size, dims, cuda, use_tqdm=use_tqdm) fid_value = calculate_frechet_distance(m1, s1, m2, s2) return fid_value
def main(cfgs): fluid.enable_imperative() if cfgs.config_str is not None: assert 'super' in cfgs.netG or 'sub' in cfgs.netG config = decode_config(cfgs.config_str) else: assert 'super' not in cfgs.model config = None data_loader, id2name = create_eval_data(cfgs, direction=cfgs.direction) model = TestModel(cfgs) model.setup() ### load_network fakes, names = [], [] for i, data in enumerate(data_loader()): model.set_input(data) if i == 0 and cfgs.need_profile: flops, params = model.profile(config) print('FLOPs: %.3fG, params: %.3fM' % (flops / 1e9, params / 1e6)) sys.exit(0) model.test(config) generated = model.fake_B fakes.append(generated.detach().numpy()) name = id2name[i] print(name) save_path = os.path.join(cfgs.save_dir, 'test') if not os.path.exists(save_path): os.makedirs(save_path) save_path = os.path.join(save_path, name) names.append(name) if i < cfgs.num_test: image = util.tensor2img(generated) util.save_image(image, save_path) fluid.disable_imperative() if not cfgs.no_fid: print('Calculating FID...') block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[2048] inception_model = InceptionV3([block_idx]) npz = np.load(cfgs.real_stat_path) fid = get_fid(fakes, inception_model, npz, cfgs.inception_model_path) print('fid score: %#.2f' % fid)
def calculate_fid_given_paths(paths, batch_size, cuda, dims): """Calculates the FID of two paths""" for p in paths: if not os.path.exists(p): raise RuntimeError('Invalid path: %s' % p) block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[dims] model = InceptionV3([block_idx]) if cuda: model.cuda() m1, s1 = _compute_statistics_of_path(paths[0], model, batch_size, dims, cuda) m2, s2 = _compute_statistics_of_path(paths[1], model, batch_size, dims, cuda) fid_value = calculate_frechet_distance(m1, s1, m2, s2) return fid_value
def test_pix2pix_fid(model, opt): opt.phase = 'val' # 根据自己的需要修改 opt.num_threads = 0 opt.batch_size = 1 opt.serial_batches = True opt.no_flip = True opt.load_size = 256 opt.display_id = -1 dataset = create_dataset(opt) model.model_eval() result_dir = os.path.join(opt.checkpoints_dir, opt.name, 'test_results') util.mkdirs(result_dir) fake_B = {} for i, data in enumerate(dataset): model.set_input(data) with torch.no_grad(): model.forward() visuals = model.get_current_visuals() fake_B[data['A_paths'][0]] = visuals['fake_B'] util.save_images(visuals, model.image_paths, result_dir, direction=opt.direction, aspect_ratio=opt.aspect_ratio) block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[2048] inception_model = InceptionV3([block_idx]) inception_model.to(model.device) inception_model.eval() # npz = np.load(os.path.join(opt.dataroot, 'real_stat_B.npz')) # fid = get_fid(list(fake_B.values()), inception_model, npz, model.device, opt.batch_size) # commented by WH at 18:29 of 2021-04-10 fid = calculate_fid_given_paths((opt.src_pics_path, opt.dst_pics_path), 1, True, 2048) return fid
def main(opt): dataloader = create_dataset(opt) device = torch.device('cuda:{}'.format(opt.gpu_ids[0])) if opt.gpu_ids \ else torch.device('cpu') block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[2048] inception_model = InceptionV3([block_idx]) inception_model.to(device) inception_model.eval() tensors = {} for i, data_i in enumerate(dataloader): tensor = data_i['B' if opt.direction == 'AtoB' else 'A'] tensors[data_i['B_paths' if opt.direction == 'AtoB' else 'A_paths'] [0]] = tensor tensors = torch.cat(list(tensors.values()), dim=0) tensors = util.tensor2imgs(tensors).astype(float) mu, sigma = _compute_statistics_of_ims(tensors, inception_model, 32, 2048, device, use_tqdm=True) np.savez(opt.output_path, mu=mu, sigma=sigma)
def __init__(self, opt): """Initialize the pix2pix class. Parameters: opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions """ assert opt.isTrain BaseModel.__init__(self, opt) # specify the training losses you want to print out. The training/test scripts will call <BaseModel.get_current_losses> self.loss_names = ['G_gan', 'G_recon', 'D_real', 'D_fake'] # specify the images you want to save/display. The training/test scripts will call <BaseModel.get_current_visuals> self.visual_names = ['real_A', 'fake_B', 'real_B'] # specify the models you want to save to the disk. The training/test scripts will call <BaseModel.save_networks> and <BaseModel.load_networks> self.model_names = ['G', 'D'] # define networks (both generator and discriminator) self.netG = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netG, opt.norm, opt.dropout_rate, opt.init_type, opt.init_gain, self.gpu_ids, opt=opt) 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) # define loss functions 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() else: raise NotImplementedError( 'Unknown reconstruction loss type [%s]!' % opt.loss_type) # initialize optimizers; schedulers will be automatically created by function <BaseModel.setup>. self.optimizer_G = torch.optim.Adam(self.netG.parameters(), 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 = [] self.optimizers.append(self.optimizer_G) self.optimizers.append(self.optimizer_D) self.eval_dataloader = create_eval_dataloader(self.opt) 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.best_fid = 1e9 self.best_mIoU = -1e9 self.fids, self.mIoUs = [], [] self.is_best = False self.Tacts, self.Sacts = {}, {} self.npz = np.load(opt.real_stat_path)
parser.add_argument('--image_size', type=int, default=256, help='image size') parser.add_argument('--batch_size', type=int, default=50, help='batch size') args = parser.parse_args() random.seed(args.seed) torch.manual_seed(args.seed) torch.cuda.manual_seed_all(args.seed) inception = InceptionV3().cuda() transform = transforms.Compose([ transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True), ]) dataset = MultiResolutionDataset(f'./dataset/{args.dataset}_lmdb', transform) loader = sample_data(dataset, args.batch_size, args.image_size) pbar = tqdm(total=len(dataset)) acts = [] for real_index, real_image in loader:
def __init__(self, opt): """Initialize the CycleGAN class. Parameters: opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions """ assert opt.isTrain assert opt.direction == 'AtoB' assert opt.dataset_mode == 'unaligned' BaseModel.__init__(self, opt) self.loss_names = [ 'D_A', 'G_A', 'G_cycle_A', 'G_idt_A', 'D_B', 'G_B', 'G_cycle_B', 'G_idt_B' ] visual_names_A = ['real_A', 'fake_B', 'rec_A'] visual_names_B = ['real_B', 'fake_A', 'rec_B'] if self.opt.lambda_identity > 0.0: visual_names_A.append('idt_B') visual_names_B.append('idt_A') self.visual_names = visual_names_A + visual_names_B self.model_names = ['G_A', 'G_B', 'D_A', 'D_B'] self.netG_A = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netG, opt.norm, opt.dropout_rate, opt.init_type, opt.init_gain, self.gpu_ids, opt=opt) self.netG_B = networks.define_G(opt.output_nc, opt.input_nc, opt.ngf, opt.netG, opt.norm, opt.dropout_rate, opt.init_type, opt.init_gain, self.gpu_ids, opt=opt) self.netD_A = 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, opt=opt) self.netD_B = networks.define_D(opt.input_nc, opt.ndf, opt.netD, opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids, opt=opt) if opt.lambda_identity > 0.0: assert (opt.input_nc == opt.output_nc) self.fake_A_pool = ImagePool(opt.pool_size) self.fake_B_pool = ImagePool(opt.pool_size) self.criterionGAN = GANLoss(opt.gan_mode).to(self.device) self.criterionCycle = torch.nn.L1Loss() self.criterionIdt = torch.nn.L1Loss() self.optimizer_G = torch.optim.Adam(itertools.chain( self.netG_A.parameters(), self.netG_B.parameters()), lr=opt.lr, betas=(opt.beta1, 0.999)) self.optimizer_D = torch.optim.Adam(itertools.chain( self.netD_A.parameters(), self.netD_B.parameters()), lr=opt.lr, betas=(opt.beta1, 0.999)) self.optimizers = [] self.optimizers.append(self.optimizer_G) self.optimizers.append(self.optimizer_D) self.eval_dataloader_AtoB = create_eval_dataloader(self.opt, direction='AtoB') self.eval_dataloader_BtoA = create_eval_dataloader(self.opt, direction='BtoA') block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[2048] self.inception_model = InceptionV3([block_idx]) self.inception_model.to(self.device) self.inception_model.eval() self.best_fid_A, self.best_fid_B = 1e9, 1e9 self.best_mIoU = -1e9 self.fids_A, self.fids_B = [], [] self.mIoUs = [] self.is_best_A = False self.is_best_B = False self.npz_A = np.load(opt.real_stat_A_path) self.npz_B = np.load(opt.real_stat_B_path)
### set configs ### if args.sched: args.lr = {128: 0.0015, 256: 0.002, 512: 0.003, 1024: 0.003} args.batch = {4: 128, 8: 64, 16: 32, 32: 16, 64: 8, 128: 8, 256: 8} else: args.lr = {} args.batch = {} args.gen_sample = {512: (8, 4), 1024: (4, 2)} args.batch_default = 8 ### prepare evaluation metrics ### inception = nn.DataParallel(InceptionV3()).cuda() gen_i, gen_j = args.gen_sample.get(args.image_size, (10, 5)) fixed_noise = torch.randn(gen_i, gen_j, code_size) real_images = torch.stack([dataset[i][1] for i in range(len(dataset))], dim=0) with open(f'./dataset/{args.dataset}_acts.pickle', 'rb') as handle: real_acts = pickle.load(handle) ### run experiments ### if args.supervised: criterion = AdaBIGGANLoss( scale_per=0.1, scale_emd=0.1,
def __init__(self, opt): """Initialize the CycleGAN class. Parameters: opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions """ assert opt.isTrain assert opt.direction == 'AtoB' assert opt.dataset_mode == 'unaligned' BaseModel.__init__(self, opt) # specify the training losses you want to print out. The training/test scripts will call <BaseModel.get_current_losses> self.loss_names = [ 'D_A', 'G_A', 'G_cycle_A', 'G_idt_A', 'D_B', 'G_B', 'G_cycle_B', 'G_idt_B' ] # specify the images you want to save/display. The training/test scripts will call <BaseModel.get_current_visuals> visual_names_A = ['real_A', 'fake_B', 'rec_A'] visual_names_B = ['real_B', 'fake_A', 'rec_B'] if self.opt.lambda_identity > 0.0: # if identity loss is used, we also visualize idt_B=G_A(B) ad idt_A=G_A(B) visual_names_A.append('idt_B') visual_names_B.append('idt_A') self.visual_names = visual_names_A + visual_names_B # combine visualizations for A and B # specify the models you want to save to the disk. The training/test scripts will call <BaseModel.save_networks> and <BaseModel.load_networks>. self.model_names = ['G_A', 'G_B', 'D_A', 'D_B'] # define networks (both Generators and discriminators) # The naming is different from those used in the paper. # Code (vs. paper): G_A (G), G_B (F), D_A (D_Y), D_B (D_X) self.netG_A = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netG, opt.norm, opt.dropout_rate, opt.init_type, opt.init_gain, self.gpu_ids) self.netG_B = networks.define_G(opt.output_nc, opt.input_nc, opt.ngf, opt.netG, opt.norm, opt.dropout_rate, opt.init_type, opt.init_gain, self.gpu_ids) self.netD_A = 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) self.netD_B = networks.define_D(opt.input_nc, opt.ndf, opt.netD, opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids) if opt.lambda_identity > 0.0: # only works when input and output images have the same number of channels assert (opt.input_nc == opt.output_nc) self.fake_A_pool = ImagePool( opt.pool_size ) # create image buffer to store previously generated images self.fake_B_pool = ImagePool( opt.pool_size ) # create image buffer to store previously generated images # define loss functions self.criterionGAN = models.modules.loss.GANLoss(opt.gan_mode).to( self.device) # define GAN loss. self.criterionCycle = torch.nn.L1Loss() self.criterionIdt = torch.nn.L1Loss() # initialize optimizers; schedulers will be automatically created by function <BaseModel.setup>. self.optimizer_G = torch.optim.Adam(itertools.chain( self.netG_A.parameters(), self.netG_B.parameters()), lr=opt.lr, betas=(opt.beta1, 0.999)) self.optimizer_D = torch.optim.Adam(itertools.chain( self.netD_A.parameters(), self.netD_B.parameters()), lr=opt.lr, betas=(opt.beta1, 0.999)) self.optimizers = [] self.optimizers.append(self.optimizer_G) self.optimizers.append(self.optimizer_D) self.eval_dataloader_AtoB = create_eval_dataloader(self.opt, direction='AtoB') self.eval_dataloader_BtoA = create_eval_dataloader(self.opt, direction='BtoA') 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 = nn.DataParallel(self.drn_model, opt.gpu_ids) self.drn_model.eval() self.best_fid_A, self.best_fid_B = 1e9, 1e9 self.best_mIoU = -1e9 self.fids_A, self.fids_B = [], [] self.mIoUs = [] self.is_best = False self.npz_A = np.load(opt.real_stat_A_path) self.npz_B = np.load(opt.real_stat_B_path)
elif 'spade' in opt.netG: # TODO raise NotImplementedError else: raise NotImplementedError configs = get_configs(config_name=opt.config_set) configs = list(configs.all_configs()) dataloader = create_dataloader(opt) model = create_model(opt) model.setup(opt) device = model.device if not opt.no_fid: block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[2048] inception_model = InceptionV3([block_idx]) inception_model.to(device) inception_model.eval() if 'cityscapes' in opt.dataroot and opt.direction == 'BtoA': drn_model = DRNSeg('drn_d_105', 19, pretrained=False) util.load_network(drn_model, opt.drn_path, verbose=False) if len(opt.gpu_ids) > 0: drn_model = nn.DataParallel(drn_model, opt.gpu_ids) drn_model.eval() npz = np.load(opt.real_stat_path) results = [] for data_i in dataloader: model.set_input(data_i) break
def __init__(self, cfgs): super(CycleGAN, self).__init__() assert cfgs.direction == 'AtoB' self.cfgs = cfgs self.loss_names = [ 'D_A', 'G_A', 'G_cycle_A', 'G_idt_A', 'D_B', 'G_B', 'G_cycle_B', 'G_idt_B' ] self.model_names = ['netG_A', 'netG_B', 'netD_A', 'netD_B'] self.netG_A = network.define_G(cfgs.input_nc, cfgs.output_nc, cfgs.ngf, cfgs.netG, cfgs.norm_type, cfgs.dropout_rate) self.netG_B = network.define_G(cfgs.output_nc, cfgs.input_nc, cfgs.ngf, cfgs.netG, cfgs.norm_type, cfgs.dropout_rate) self.netD_A = network.define_D(cfgs.output_nc, cfgs.ndf, cfgs.netD, cfgs.norm_type, cfgs.n_layer_D) self.netD_B = network.define_D(cfgs.input_nc, cfgs.ndf, cfgs.netD, cfgs.norm_type, cfgs.n_layer_D) if self.cfgs.use_parallel: self.netG_A = fluid.dygraph.parallel.DataParallel( self.netG_A, self.cfgs.strategy) self.netG_B = fluid.dygraph.parallel.DataParallel( self.netG_B, self.cfgs.strategy) self.netD_A = fluid.dygraph.parallel.DataParallel( self.netD_A, self.cfgs.strategy) self.netD_B = fluid.dygraph.parallel.DataParallel( self.netD_B, self.cfgs.strategy) if cfgs.lambda_identity > 0.0: assert (cfgs.input_nc == cfgs.output_nc) self.fake_A_pool = ImagePool(cfgs.pool_size) self.fake_B_pool = ImagePool(cfgs.pool_size) self.optimizer_G = optimization.Optimizer( cfgs.mobile_lr, cfgs.mobile_scheduler, cfgs.step_per_epoch, cfgs.mobile_nepochs, cfgs.mobile_nepochs_decay, cfgs, parameter_list=(self.netG_A.parameters() + self.netG_B.parameters())) self.optimizer_D_A = optimization.Optimizer( cfgs.mobile_lr, cfgs.mobile_scheduler, cfgs.step_per_epoch, cfgs.mobile_nepochs, cfgs.mobile_nepochs_decay, cfgs, parameter_list=self.netD_A.parameters()) self.optimizer_D_B = optimization.Optimizer( cfgs.mobile_lr, cfgs.mobile_scheduler, cfgs.step_per_epoch, cfgs.mobile_nepochs, cfgs.mobile_nepochs_decay, cfgs, parameter_list=self.netD_B.parameters()) self.eval_dataloader_AtoB, self.name_AtoB = create_eval_data( cfgs, direction='AtoB') self.eval_dataloader_BtoA, self.name_BtoA = create_eval_data( cfgs, direction='BtoA') block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[2048] self.inception_model = InceptionV3([block_idx]) self.best_fid_A, self.best_fid_B = 1e9, 1e9 self.fids_A, self.fids_B = [], [] self.is_best = False self.npz_A = np.load(cfgs.real_stat_A_path) self.npz_B = np.load(cfgs.real_stat_B_path)
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, opt): """Initialize the pix2pix class. Parameters: opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions """ assert opt.isTrain BaseModel.__init__(self, opt) self.loss_names = [ 'G_gan', 'G_recon', 'D_real', 'D_fake', 'G_comp_cost' ] self.visual_names = ['real_A', 'fake_B', 'real_B'] self.model_names = ['G', 'D'] self.netG = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netG, opt.norm, opt.dropout_rate, opt.init_type, opt.init_gain, self.gpu_ids, opt=opt) 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, opt=opt) 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() else: raise NotImplementedError( 'Unknown reconstruction loss type [%s]!' % opt.loss_type) self.optimizer_G = torch.optim.Adam(self.netG.parameters(), 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 = [] self.optimizers.append(self.optimizer_G) self.optimizers.append(self.optimizer_D) self.eval_dataloader = create_eval_dataloader(self.opt) 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.best_fid = 1e9 self.best_mIoU = -1e9 self.fids, self.mIoUs = [], [] self.is_best = False self.Tacts, self.Sacts = {}, {} self.npz = np.load(opt.real_stat_path)
def __init__(self, cfgs, task): super(BaseResnetDistiller, self).__init__() self.cfgs = cfgs self.task = task self.loss_names = ['G_gan', 'G_distill', 'G_recon', 'D_fake', 'D_real'] self.model_names = ['netG_student', 'netG_teacher', 'netD'] self.netG_teacher = network.define_G(cfgs.input_nc, cfgs.output_nc, cfgs.ngf, cfgs.netG, cfgs.norm_type, cfgs.dropout_rate) student_netG = getattr(cfgs, self.task + '_student_netG') self.netG_student = network.define_G(cfgs.input_nc, cfgs.output_nc, cfgs.student_ngf, student_netG, cfgs.norm_type, cfgs.student_dropout_rate) if self.task == 'distiller': self.netG_pretrained = network.define_G(cfgs.input_nc, cfgs.output_nc, cfgs.pretrained_ngf, cfgs.pretrained_netG, cfgs.norm_type, 0) if self.cfgs.use_parallel: self.netG_pretrained = fluid.dygraph.parallel.DataParallel( self.netG_pretrained, self.cfgs.strategy) self.netD = network.define_D(cfgs.output_nc, cfgs.ndf, cfgs.netD, cfgs.norm_type, cfgs.n_layer_D) if self.cfgs.use_parallel: self.netG_teacher = fluid.dygraph.parallel.DataParallel( self.netG_teacher, self.cfgs.strategy) self.netG_student = fluid.dygraph.parallel.DataParallel( self.netG_student, self.cfgs.strategy) self.netD = fluid.dygraph.parallel.DataParallel( self.netD, self.cfgs.strategy) self.netG_teacher.eval() self.netG_student.train() self.netD.train() ### [9, 12, 15, 18] self.mapping_layers = [ '_layers.model.%d' % i for i in range(9, 21, 3) ] if self.cfgs.use_parallel else [ '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 = cfgs.ngf, cfgs.student_ngf if self.task == 'distiller': netA = Conv2D(num_channels=fs * 4, num_filters=ft * 4, filter_size=1) else: netA = SuperConv2D(num_channels=fs * 4, num_filters=ft * 4, filter_size=1) G_params += netA.parameters() self.netAs.append(netA) self.loss_names.append('G_distill%d' % i) if self.task == 'distiller': learning_rate = cfgs.distiller_lr scheduler = cfgs.distiller_scheduler nepochs = cfgs.distiller_nepochs nepochs_decay = cfgs.distiller_nepochs_decay elif self.task == 'supernet': learning_rate = cfgs.supernet_lr scheduler = cfgs.supernet_scheduler nepochs = cfgs.supernet_nepochs nepochs_decay = cfgs.supernet_nepochs_decay else: raise NotImplementedError("task {} is not suppport".format( self.task)) self.optimizer_G = optimization.Optimizer(learning_rate, scheduler, cfgs.step_per_epoch, nepochs, nepochs_decay, cfgs, parameter_list=G_params) self.optimizer_D = optimization.Optimizer( learning_rate, scheduler, cfgs.step_per_epoch, nepochs, nepochs_decay, cfgs, parameter_list=self.netD.parameters()) self.eval_dataloader, self.name = create_eval_data( cfgs, direction=cfgs.direction) block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[2048] self.inception_model = InceptionV3([block_idx]) self.is_best = False if cfgs.real_stat_path: self.npz = np.load(cfgs.real_stat_path) self.is_best = False