def load_checkpoints(config, checkpoint, blend_scale=0.125, first_order_motion_model=False, cpu=False):
with open(config) as f:
config = yaml.load(f)
reconstruction_module = PartSwapGenerator(blend_scale=blend_scale,
first_order_motion_model=first_order_motion_model,
**config['model_params']['reconstruction_module_params'],
**config['model_params']['common_params'])
if not cpu:
reconstruction_module.cuda()
segmentation_module = SegmentationModule(**config['model_params']['segmentation_module_params'],
**config['model_params']['common_params'])
if not cpu:
segmentation_module.cuda()
if cpu:
checkpoint = torch.load(checkpoint, map_location=torch.device('cpu'))
else:
checkpoint = torch.load(checkpoint)
load_reconstruction_module(reconstruction_module, checkpoint)
load_segmentation_module(segmentation_module, checkpoint)
if not cpu:
reconstruction_module = DataParallelWithCallback(reconstruction_module)
segmentation_module = DataParallelWithCallback(segmentation_module)
reconstruction_module.eval()
segmentation_module.eval()
return reconstruction_module, segmentation_module
with open(opt.config) as f:
config = yaml.load(f)
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())
reconstruction_module = ReconstructionModule(
**config['model_params']['reconstruction_module_params'],
**config['model_params']['common_params'])
reconstruction_module.to(opt.device_ids[0])
if opt.verbose:
print(reconstruction_module)
segmentation_module = SegmentationModule(
**config['model_params']['segmentation_module_params'],
**config['model_params']['common_params'])
segmentation_module.to(opt.device_ids[0])
if opt.verbose:
print(segmentation_module)
dataset = FramesDataset(is_train=True, **config['dataset_params'])
if not os.path.exists(log_dir):
os.makedirs(log_dir)
if not os.path.exists(os.path.join(log_dir, os.path.basename(opt.config))):
copy(opt.config, log_dir)
print("Training...")
train(config, reconstruction_module, segmentation_module, opt.checkpoint,
log_dir, dataset, opt.device_ids)