def load_checkpoints(config_path, checkpoint_path, cpu=False):
with open(config_path) as f:
config = yaml.load(f)
generator = OcclusionAwareGenerator(**config['model_params']['generator_params'],
**config['model_params']['common_params'])
if not cpu:
generator.cuda()
kp_detector = KPDetector(**config['model_params']['kp_detector_params'],
**config['model_params']['common_params'])
if not cpu:
kp_detector.cuda()
if cpu:
checkpoint = torch.load(checkpoint_path, map_location=torch.device('cpu'))
else:
checkpoint = torch.load(checkpoint_path)
generator.load_state_dict(checkpoint['generator'])
kp_detector.load_state_dict(checkpoint['kp_detector'])
if not cpu:
generator = DataParallelWithCallback(generator)
kp_detector = DataParallelWithCallback(kp_detector)
generator.eval()
kp_detector.eval()
return generator, kp_detector
def load_checkpoints(config_path, checkpoint_path, device='cuda'):
with open(config_path) as f:
config = yaml.load(f)
generator = OcclusionAwareGenerator(
**config['model_params']['generator_params'],
**config['model_params']['common_params'])
generator.to(device)
kp_detector = KPDetector(**config['model_params']['kp_detector_params'],
**config['model_params']['common_params'])
kp_detector.to(device)
checkpoint = torch.load(checkpoint_path, map_location=device)
generator.load_state_dict(checkpoint['generator'])
kp_detector.load_state_dict(checkpoint['kp_detector'])
generator = DataParallelWithCallback(generator)
kp_detector = DataParallelWithCallback(kp_detector)
generator.eval()
kp_detector.eval()
return generator, kp_detector
def load_checkpoints(config_path, checkpoint_path, device="cuda"):
with open(config_path) as f:
config = yaml.load(f, Loader=yaml.FullLoader)
generator = OcclusionAwareGenerator(
**config["model_params"]["generator_params"],
**config["model_params"]["common_params"],
)
generator.to(device)
kp_detector = KPDetector(
**config["model_params"]["kp_detector_params"],
**config["model_params"]["common_params"],
)
kp_detector.to(device)
checkpoint = torch.load(checkpoint_path, map_location=device)
generator.load_state_dict(checkpoint["generator"])
kp_detector.load_state_dict(checkpoint["kp_detector"])
generator = DataParallelWithCallback(generator)
kp_detector = DataParallelWithCallback(kp_detector)
generator.eval()
kp_detector.eval()
return generator, kp_detector
def load_checkpoints(config_path):
with open(config_path) as f:
config = yaml.load(f)
pretrain_model = config['ckpt_model']
generator = OcclusionAwareGenerator(
**config['model_params']['generator_params'],
**config['model_params']['common_params'])
kp_detector = KPDetector(**config['model_params']['kp_detector_params'],
**config['model_params']['common_params'])
load_ckpt(pretrain_model, generator=generator, kp_detector=kp_detector)
generator.eval()
kp_detector.eval()
return generator, kp_detector
def load_checkpoints(config_path):
with open(config_path) as f:
config = yaml.load(f)
pretrain_model = config['ckpt_model']
generator = OcclusionAwareGenerator(
**config['model_params']['generator_params'],
**config['model_params']['common_params'])
kp_detector = KPDetector(**config['model_params']['kp_detector_params'],
**config['model_params']['common_params'])
if pretrain_model['generator'] is not None:
if pretrain_model['generator'][-3:] == 'npz':
G_param = np.load(pretrain_model['generator'],
allow_pickle=True)['arr_0'].item()
G_param_clean = [(i, G_param[i]) for i in G_param
if 'num_batches_tracked' not in i]
parameter_clean = generator.parameters()
del (
parameter_clean[65]
) # The parameters in AntiAliasInterpolation2d is not in dict_set and should be ignore.
for p, v in zip(parameter_clean, G_param_clean):
p.set_value(v[1])
else:
a, b = fluid.load_dygraph(pretrain_model['generator'])
generator.set_dict(a)
print('Restore Pre-trained Generator')
if pretrain_model['kp'] is not None:
if pretrain_model['kp'][-3:] == 'npz':
KD_param = np.load(pretrain_model['kp'],
allow_pickle=True)['arr_0'].item()
KD_param_clean = [(i, KD_param[i]) for i in KD_param
if 'num_batches_tracked' not in i]
parameter_clean = kp_detector.parameters()
for p, v in zip(parameter_clean, KD_param_clean):
p.set_value(v[1])
else:
a, b = fluid.load_dygraph(pretrain_model['kp'])
kp_detector.set_dict(a)
print('Restore Pre-trained KD')
generator.eval()
kp_detector.eval()
return generator, kp_detector
# read in the config file
config = yaml.load(
f
) # config file contains code directions, including training details
# generate a log path (store running time details)
if opt.checkpoint is not None:
log_dir = os.path.join(*os.path.split(opt.checkpoint)[:-1])
else:
log_dir = os.path.join(opt.log_dir,
os.path.basename(opt.config).split('.')[0])
log_dir += ' ' + strftime("%d_%m_%y_%H.%M.%S", gmtime())
# Declare an image generator
generator = OcclusionAwareGenerator(
**config['model_params']['generator_params'],
**config['model_params']['common_params'])
# If GPU Available, adapt to it
if torch.cuda.is_available():
generator.to(opt.device_ids[0])
if opt.verbose:
print(generator)
# Declare a discriminator
discriminator = MultiScaleDiscriminator(
**config['model_params']['discriminator_params'],
**config['model_params']['common_params'])
if torch.cuda.is_available():
discriminator.to(opt.device_ids[0])
if opt.verbose:
def load_checkpoints(config_path, checkpoint_path, cpu=False):
with open(config_path) as f:
config = yaml.load(f)
generator = OcclusionAwareGenerator(
**config["model_params"]["generator_params"],
**config["model_params"]["common_params"],
)
if cpu:
generator.cpu()
else:
generator.cuda()
kp_detector = KPDetector(
**config["model_params"]["kp_detector_params"],
**config["model_params"]["common_params"],
)
if cpu:
kp_detector.cpu()
else:
kp_detector.cuda()
checkpoint = torch.load(checkpoint_path, map_location="cpu" if cpu else None)
generator.load_state_dict(checkpoint["generator"])
kp_detector.load_state_dict(checkpoint["kp_detector"])
generator = DataParallelWithCallback(generator)
kp_detector = DataParallelWithCallback(kp_detector)
generator.eval()
kp_detector.eval()
return generator, kp_detector
parser.add_argument("--mode",
default="train",
choices=["train", "reconstruction", "animate"])
parser.add_argument("--save_dir",
default='/home/aistudio/train_ckpt',
help="path to save in")
parser.add_argument("--preload",
action='store_true',
help="preload dataset to RAM")
parser.set_defaults(verbose=False)
opt = parser.parse_args()
with open(opt.config) as f:
config = yaml.load(f)
generator = OcclusionAwareGenerator(
**config['model_params']['generator_params'],
**config['model_params']['common_params'])
discriminator = MultiScaleDiscriminator(
**config['model_params']['discriminator_params'],
**config['model_params']['common_params'])
kp_detector = KPDetector(**config['model_params']['kp_detector_params'],
**config['model_params']['common_params'])
dataset = FramesDataset(is_train=(opt.mode == 'train'),
**config['dataset_params'])
if opt.preload:
logging.info('PreLoad Dataset: Start')
pre_list = list(range(len(dataset)))
import multiprocessing.pool as pool
with pool.Pool(4) as pl:
buf = pl.map(dataset.preload, pre_list)
