def ng_update(self, variables, loss=None, inverted_loss=False):
"""
Updates NG distribution either with the provided loss or loss that
is recomputed.
Args:
variables (dict): a dictionary instance generated from the
variable manager.
loss (array or list): a 1-dimensional array or list consisting of
losses corresponding to each sample. If the loss is not
provided, uses the variables to recompute the loss.
[Default: None]
inverted_loss (bool): if True, the loss is computed after inverting
the generated images back to the original target. For example
this is used to compute the loss on the original target.
[Default: False]
"""
for (var_type, var_name), ng_opt in self.ng_optimizers.items():
ng_data = self._sampled[(var_type, var_name)]
if loss is None:
out, loss, _ = self.step(variables, optimize=False)
if inverted_loss and hasattr(variables, 'transform'):
target_type = \
self.var_manager.variable_info['target']['var_type']
weight_type = \
self.var_manager.variable_info['weight']['var_type']
target = self.var_manager.variable_info['target']['default']
weight = self.var_manager.variable_info['weight']['default']
target = target.unsqueeze(0).type_as(out)
weight = weight.unsqueeze(0).type_as(out)
t_fn = self.transform_fns['target']['fn']
t_param = torch.stack(variables.transform.t.data)
out = t_fn(out, t_param, invert=True)
loss = self.loss_fn(out, target, binarize(weight))
loss = loss.cpu().detach().numpy()
for d, l in zip(ng_data, loss):
ng_opt.tell(d, l)
return
### ---- 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)
def compute_pre_alignment(weight):
""" Precompute intialization based on BigGAN bias """
dst_center, dst_size = get_biggan_stats()
src_center, src_size = compute_stat_from_mask(binarize(weight))
t = convert_to_t(src_center, src_size, dst_center, dst_size)
return t.numpy()