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()