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