opt.isTrain = False
if opt.load_path is None or not os.path.exists(opt.load_path):
raise FileExistsError('Load path must be exist!!!')
device = torch.device(f'cuda:{opt.gpu_ids[0]}') if len(
opt.gpu_ids) > 0 else 'cpu'
ckpt = torch.load(opt.load_path, map_location=device)
cfg = ckpt['cfg'] if 'cfg' in ckpt.keys() else (None, None)
# create model
if opt.model == 'cyclegan':
if opt.mask:
model = MaskCycleGAN.MaskCycleGANModel(opt)
else:
model = CycleGAN.CycleGANModel(opt,
cfg_AtoB=cfg[0],
cfg_BtoA=cfg[1])
elif opt.model == 'pix2pix':
opt.norm = 'batch'
opt.dataset_mode = 'aligned'
opt.pool_size = 0
if opt.mask:
model = MaskPix2Pix.MaskPix2PixModel(opt)
else:
model = Pix2Pix.Pix2PixModel(opt,
filter_cfgs=cfg[0],
channel_cfgs=cfg[1])
elif opt.model == 'mobilecyclegan':
if opt.mask:
model = MaskMobileCycleGAN.MaskMobileCycleGANModel(opt)
else:
opt.isTrain = True
util.mkdirs(os.path.join(opt.checkpoints_dir, opt.name))
logger = util.get_logger(
os.path.join(opt.checkpoints_dir, opt.name, 'logger.log'))
best_AtoB_fid = float('inf') if 'cityscapes' not in opt.dataroot else 0.0
best_BtoA_fid = float('inf') if 'cityscapes' not in opt.dataroot else 0.0
best_AtoB_epoch = 0
best_BtoA_epoch = 0
# create model
if opt.model == 'cyclegan':
if opt.mask:
model = MaskCycleGAN.MaskCycleGANModel(opt)
else:
model = CycleGAN.CycleGANModel(opt)
elif opt.model == 'pix2pix':
opt.norm = 'batch'
opt.dataset_mode = 'aligned'
opt.pool_size = 0
if opt.mask:
model = MaskPix2Pix.MaskPix2PixModel(opt)
else:
model = Pix2Pix.Pix2PixModel(opt)
elif opt.model == 'mobilecyclegan':
if opt.mask:
model = MaskMobileCycleGAN.MaskMobileCycleGANModel(opt)
else:
model = MobileCycleGAN.MobileCycleGANModel(opt)
elif opt.model == 'mobilepix2pix':
opt.norm = 'batch'