### ---- optimize --- ###
opt = HybridNevergradOptimizer(args.ng_method,
model,
var_manager,
loss_fn,
max_batch_size=args.max_minibatch,
log=args.make_video)
vars, out, loss = opt.optimize(
num_samples=args.num_samples,
meta_steps=30,
grad_steps=50,
last_grad_steps=300,
)
### ---- save results ---- #
vars.loss = loss
os.makedirs(save_dir, exist_ok=True)
save_variables(osp.join(save_dir, 'vars.npy'), vars)
if args.make_video:
video.make_video(osp.join(save_dir, 'out.mp4'), out)
image.save(osp.join(save_dir, 'target.jpg'), target)
image.save(osp.join(save_dir, 'mask.jpg'), image.binarize(weight))
image.save(osp.join(save_dir, 'out.jpg'), out[-1])
np.save(osp.join(save_dir, 'tracked.npy'), opt.tracked)
t_opt.register_transform(weight_transform_fn, 't', 'weight')
# (highly recommended) speeds up optimization by propating information
t_opt.set_variable_propagation('z')
t_vars, (t_out, t_target, t_candidate), t_loss = \
t_opt.optimize(meta_steps=50, grad_steps=10)
os.makedirs(save_dir, exist_ok=True)
if args.make_video:
video.make_video(osp.join(save_dir, 'transform_out.mp4'), t_out)
video.make_video(osp.join(save_dir, 'transform_target.mp4'), t_target)
image.save(osp.join(save_dir, 'transform_out.jpg'), t_out[-1])
image.save(osp.join(save_dir, 'transform_target.jpg'), t_target[-1])
image.save(osp.join(save_dir, 'transform_candidate.jpg'), t_candidate)
np.save(osp.join(save_dir, 'transform_tracked.npy'),
{'t': t_opt.transform_tracked})
t = t_opt.get_candidate()
var_manager.edit_variable('t', {'default': t, 'grad_free': False})
var_manager.edit_variable('z', {'learning_rate': args.lr})
del t_opt, t_vars, t_out, t_target, t_candidate, t_loss
model.zero_grad()
torch.cuda.empty_cache()
