def get_reg(self, **kwargs): """ Get weights regularization (KL(q(w)||p(w)) approximation) """ log_alp = self.clip_alp(self.log_alp) element_wise_kl = .5 * log_alp \ + 1.16145124 * torch.exp(log_alp) \ - 1.50204118 * torch.exp(log_alp) ** 2 \ + 0.58629921 * torch.exp(log_alp) ** 3 sum_kl = element_wise_kl.sum(dim=(1, 2, 3)) beta = F.sigmoid(self.clip_beta(self.beta_r)) qz = torch.cat([self.ONE, torch.cumprod(beta, dim=0)]) * torch.cat( [1 - beta, self.ONE]) coef0 = torch.cumsum(qz, dim=0)[:-1] coef1 = torch.sum(qz) - coef0 coef1 = torch.cat([self.ONE.repeat(self.FREEZE_PART), coef1]).repeat_interleave(self.EVERY) kl_w = coef1.dot(sum_kl) qz = torch.cat([self.ONE, torch.cumprod(beta, dim=0)]) * torch.cat( [1 - beta, self.ONE]) log_frac_qz_pz = torch.log(qz / self.pz) kl_z = qz.dot(log_frac_qz_pz) kl = -(kl_w - kl_z) return kl
def cumprod(x, axis: int = 0, exclusive: bool = False): if exclusive: x = torch.transpose(x, axis, -1) x = torch.cat((torch.ones_like(x[..., -1:]), x[..., :-1]), -1) res = torch.cumprod(x, -1) return torch.transpose(res, axis, -1) return torch.cumprod(x, axis)
def gen_sampler(self): # 初始化脊椎+椎间盘标注数据list v_annos, d_annos = [], [] # 遍历所有标注数据中的key和对应的脊椎+椎间盘标注数据 for key, (v_anno, d_anno) in self.annotations: # 将对应的数据添加到相应的list中 v_annos.append(v_anno[:, -1]) d_annos.append(d_anno[:, -1]) # 将所有脊椎+椎间盘标注数据进行堆叠 v_annos = torch.stack(v_annos, dim=0) d_annos = torch.stack(d_annos, dim=0) # 统计标注数据中每个独立元素的个数 v_count = torch.unique(v_annos, return_counts=True)[1] d_count = torch.unique(d_annos, return_counts=True)[1] # 首先对独立元素个数的最后一维进行cumprod操作,获取到结果的最后一列数据后,将其除1,返回除法的浮点数结果而不作整数处理 v_weights = torch.true_divide( 1, torch.cumprod(v_count[v_annos], dim=-1)[:, -1]) d_weights = torch.true_divide( 1, torch.cumprod(d_count[d_annos], dim=-1)[:, -1]) # 根据脊椎权重和椎间盘权重计算得到总权重 weights = v_weights * d_weights # 使用计算好的权重进行加权随机采样,返回采样得到的数据 return WeightedRandomSampler(weights=weights, num_samples=len(self), replacement=True)
def fb(self, x, mode='batch'): if mode == 'batch': # TODO(j-pong): multi-target energy x = 1.0 - self.minimaxn(x[:, 0]) bsz, tsz = x.size() xs = x.unsqueeze(-1).repeat(1, 1, tsz) ret = [] for x in xs: m_f = torch.triu(torch.ones(tsz, tsz)).to(x.device) x_f = torch.cumprod(torch.tril(x) + m_f, dim=-2) - m_f m_b = torch.tril(torch.ones(tsz, tsz)).to(x.device) x_b = torch.cumprod((torch.triu(x) + m_b).flip(dims=[-2]), dim=-2).flip(dims=[-2]) - m_b ret.append(x_b + x_f) xs = torch.stack(ret, dim=0).unsqueeze(1) / 2 # B, 1, T, T else: x = 1.0 - self.minimaxn(x[:, 0:1]) xs = [] for i in range(x.size(-1)): if i != 0: x_f = torch.cumprod(x[:, :, i:], dim=-1) x_b = torch.cumprod(x[:, :, :i].flip(dims=[-1]), dim=-1).flip(dims=[-1]) xs.append(torch.cat([x_b, x_f], dim=-1)) else: xs.append(torch.cumprod(x[:, :, :], dim=-1)) xs = torch.stack(xs, dim=-1) # B, tnum, T, T return xs
def cpt_gate(self, semantic_score: T) -> Tuple[T, T]: assert semantic_score.size()[0] > 4 score = semantic_score[1:-1] # (num_score - 2) fwd_score = torch.cat([torch.zeros(score.size(0)), score], dim=0) bwd_score = torch.cat([score, torch.zeros(score.size(0))], dim=0) fwd_score_hat = torch.stack([ fwd_score[i:i + score.size()[0]] for i in range(score.size()[0] - 1, 0, -1) ], dim=0) bwd_score_hat = torch.stack([ bwd_score[i:i + score.size(0)] for i in range(1, score.size()[0]) ], dim=0) if self.hard: fwd_gate = (F.hardtanh( (fwd_score_hat - score[None, :]) / self.resolution * 2 + 1) + 1) / 2 bwd_gate = (F.hardtanh( (bwd_score_hat - score[None, :]) / self.resolution * 2 + 1) + 1) / 2 else: fwd_gate = F.sigmoid((fwd_score_hat - score[None, :]) / self.resolution * 10 + 5) bwd_gate = F.sigmoid((bwd_score_hat - score[None, :]) / self.resolution * 10 + 5) fwd_gate = torch.cumprod(fwd_gate, dim=0) # seq x seq - 1 bwd_gate = torch.cumprod(bwd_gate, dim=0) # seq x seq - 1 return (fwd_gate, bwd_gate)
def fb(x, mode='batch'): # TODO(j-pong): multi-target energy with torch.no_grad(): if mode == 'batch': x = x[:, 0] bsz, tsz = x.size() xs = x.unsqueeze(-1).repeat(1, 1, tsz) ret = [] for x in xs: ones = torch.ones(tsz, tsz).to(x.device) m_f = torch.triu(ones, diagonal=1) x_f = torch.cumprod(torch.tril(x) + m_f, dim=-2) x_f = torch.cat([ones[0:1], x_f[:-1]], dim=0) - m_f ret.append(x_f) xs = torch.stack(ret, dim=0).unsqueeze(1) # B, 1, T, T else: xs = [] for i in range(x.size(-1)): if i != 0: x_f = torch.cumprod(x[:, :, i:], dim=-1) x_b = torch.cumprod(x[:, :, :i].flip(dims=[-1]), dim=-1).flip(dims=[-1]) xs.append(torch.cat([x_b, x_f], dim=-1)) else: xs.append(torch.cumprod(x[:, :, :], dim=-1)) xs = torch.stack(xs, dim=-1) # B, tnum, T, T if torch.isnan(xs.sum()): raise ValueError("kernel target value has nan") return xs
def compute_gate(self, score: T) -> Tuple[T, T]: #assert score.size()[0] > 4 score = score[1:-1] fwd_gate = self.compute_prob(self.pad_score(score), score) bwd_gate = self.compute_prob( self.pad_score(torch.flip(score, dims=(0, ))), torch.flip(score, dims=(0, ))) fwd_gate = torch.cumprod(fwd_gate, dim=0) # seq x seq - 1 bwd_gate = torch.cumprod(bwd_gate, dim=0) # seq x seq - 1 return (fwd_gate, bwd_gate)
def raw2outputs_blending(raw_dy, raw_rigid, raw_blend_w, z_vals, rays_d, raw_noise_std): act_fn = F.relu dists = z_vals[..., 1:] - z_vals[..., :-1] dists = torch.cat( [dists, torch.Tensor([1e10]).expand(dists[..., :1].shape)], -1) # [N_rays, N_samples] dists = dists * torch.norm(rays_d[..., None, :], dim=-1) rgb_dy = torch.sigmoid(raw_dy[..., :3]) # [N_rays, N_samples, 3] rgb_rigid = torch.sigmoid(raw_rigid[..., :3]) # [N_rays, N_samples, 3] noise = 0. if raw_noise_std > 0.: noise = torch.randn(raw_dy[..., 3].shape) * raw_noise_std opacity_dy = act_fn(raw_dy[..., 3] + noise) #.detach() #* raw_blend_w opacity_rigid = act_fn(raw_rigid[..., 3] + noise) #.detach() #* (1. - raw_blend_w) # alpha with blending weights alpha_dy = (1. - torch.exp(-opacity_dy * dists)) * raw_blend_w alpha_rig = (1. - torch.exp(-opacity_rigid * dists)) * (1. - raw_blend_w) Ts = torch.cumprod( torch.cat([ torch.ones((alpha_dy.shape[0], 1)), (1. - alpha_dy) * (1. - alpha_rig) + 1e-10 ], -1), -1)[:, :-1] weights_dy = Ts * alpha_dy weights_rig = Ts * alpha_rig # union map rgb_map = torch.sum(weights_dy[..., None] * rgb_dy + \ weights_rig[..., None] * rgb_rigid, -2) weights_mix = weights_dy + weights_rig depth_map = torch.sum(weights_mix * z_vals, -1) # compute dynamic depth only alpha_dynamic = 1. - torch.exp(-opacity_dy * dists) weights_dynamic = alpha_dynamic * torch.cumprod( torch.cat([ torch.ones((alpha_dynamic.shape[0], 1)), 1. - alpha_dynamic + 1e-10 ], -1), -1)[:, :-1] depth_map_dynamic = torch.sum(weights_dynamic * z_vals, -1) rgb_map_dy = torch.sum( weights_dynamic[..., None] * torch.sigmoid(raw_dy[..., :3]), -2) return rgb_map, depth_map, \ rgb_map_dy, depth_map_dynamic, weights_dynamic
def torch_ideal_err(sorted_labels, k=10, point=True, gpu=False): assert sorted_labels.size(0) >= k max_label = torch.max(sorted_labels) labels = sorted_labels[0:k] satis_pros = (torch.pow(2.0, labels) - 1.0) / torch.pow(2.0, max_label) unsatis_pros = torch.ones_like(labels) - satis_pros cum_unsatis_pros = torch.cumprod(unsatis_pros, dim=0) if gpu: ranks = torch.arange(k).type(tensor) + 1.0 expt_ranks = 1.0 / ranks else: ranks = torch.arange(k) + 1.0 expt_ranks = 1.0 / ranks cascad_unsatis_pros = ranks cascad_unsatis_pros[1:k] = cum_unsatis_pros[0:k - 1] expt_satis_ranks = expt_ranks * satis_pros * cascad_unsatis_pros # w.r.t. all rank positions if point: # a specific position ideal_err = torch.sum(expt_satis_ranks, dim=0) return ideal_err else: ideal_err_at_ks = torch.cumsum(expt_satis_ranks, dim=0) return ideal_err_at_ks
def update_allocation_weights(self, sorted_indices, ordered_usage): ''' Update allocation weights (B, N) :param (B, R) free gates ''' ones = torch.ones(self.batch_size, device=self.device).view(-1, 1) usage_ones = torch.cat((ones, ordered_usage[:, :-1]), dim=1) prod = torch.cumprod(usage_ones, dim=1) #prod = torch.cat((ones, prod), dim=1) allocation_weights_ordered = (1 - ordered_usage) * prod # reorder allocation weights - inefficient but simpler version #for b in self.batch_size: # self.allocation_weights[b][sorted_indices[b]] = allocation_weights_ordered[b] # more efficient version than above - avoid a loop over B dimension # N.B. this part is crucial for the convergence of the Copy task allocation_weights = self.unorder_tensor(allocation_weights_ordered, sorted_indices) return allocation_weights
def _compute_loss_actor(self, imag_beliefs, imag_states, imag_ac_logps=None): # reward and value prediction of imagined trajectories imag_rewards = bottle(self.reward_model, (imag_beliefs, imag_states)) imag_values = bottle(self.value_model, (imag_beliefs, imag_states)) with torch.no_grad(): if self.args.pcont: pcont = bottle(self.pcont_model, (imag_beliefs, imag_states)) else: pcont = self.args.discount * torch.ones_like(imag_rewards) pcont = pcont.detach() if imag_ac_logps is not None: imag_values[ 1:] -= self.args.temp * imag_ac_logps # add entropy here returns = cal_returns(imag_rewards[:-1], imag_values[:-1], imag_values[-1], pcont[:-1], lambda_=self.args.disclam) discount = torch.cumprod( torch.cat([torch.ones_like(pcont[:1]), pcont[:-2]], 0), 0).detach() actor_loss = -torch.mean(discount * returns) return actor_loss
def cumprod_exclusive(tensor: torch.Tensor) -> torch.Tensor: r"""Mimick functionality of tf.math.cumprod(..., exclusive=True), as it isn't available in PyTorch. Args: tensor (torch.Tensor): Tensor whose cumprod (cumulative product, see `torch.cumprod`) along dim=-1 is to be computed. Returns: cumprod (torch.Tensor): cumprod of Tensor along dim=-1, mimiciking the functionality of tf.math.cumprod(..., exclusive=True) (see `tf.math.cumprod` for details). """ # TESTED # Only works for the last dimension (dim=-1) dim = -1 # Compute regular cumprod first (this is equivalent to `tf.math.cumprod(..., exclusive=False)`). cumprod = torch.cumprod(tensor, dim) # "Roll" the elements along dimension 'dim' by 1 element. cumprod = torch.roll(cumprod, 1, dim) # Replace the first element by "1" as this is what tf.cumprod(..., exclusive=True) does. cumprod[..., 0] = 1.0 return cumprod
def torch_batch_ideal_err(batch_sorted_labels, k=10, gpu=False, point=True): assert batch_sorted_labels.size(1) > k batch_max = torch.max(batch_sorted_labels, dim=1) batch_labels = batch_sorted_labels[:, 0:k] batch_satis_pros = (torch.pow(2.0, batch_labels) - 1.0) / torch.pow( 2.0, batch_max) batch_unsatis_pros = torch.ones(batch_labels) - batch_satis_pros batch_cum_unsatis_pros = torch.cumprod(batch_unsatis_pros, dim=1) positions = torch.arange(k) + 1.0 positions = positions.view(1, -1) positions = torch.repeat_interleave(positions, batch_sorted_labels.size(0), dim=0) batch_expt_ranks = 1.0 / positions cascad_unsatis_pros = positions cascad_unsatis_pros[:, 1:k] = batch_cum_unsatis_pros[:, 0:k - 1] expt_satis_ranks = batch_expt_ranks * batch_satis_pros * cascad_unsatis_pros # w.r.t. all rank positions if point: batch_errs = torch.sum(expt_satis_ranks, dim=1) return batch_errs else: batch_err_at_ks = torch.cumsum(expt_satis_ranks, dim=1) return batch_err_at_ks
def __init__(self, betas, model_mean_type, model_var_type, loss_type): super().__init__() betas = betas.type(torch.float64) timesteps = betas.shape[0] self.num_timesteps = int(timesteps) self.model_mean_type = model_mean_type # xprev, xstart, eps self.model_var_type = model_var_type # learned, fixedsmall, fixedlarge self.loss_type = loss_type # kl, mse alphas = 1 - betas alphas_cumprod = torch.cumprod(alphas, 0) alphas_cumprod_prev = torch.cat( (torch.tensor([1], dtype=torch.float64), alphas_cumprod[:-1]), 0 ) posterior_variance = betas * (1 - alphas_cumprod_prev) / (1 - alphas_cumprod) self.register("betas", betas) self.register("alphas_cumprod", alphas_cumprod) self.register("alphas_cumprod_prev", alphas_cumprod_prev) self.register("sqrt_alphas_cumprod", torch.sqrt(alphas_cumprod)) self.register("sqrt_one_minus_alphas_cumprod", torch.sqrt(1 - alphas_cumprod)) self.register("log_one_minus_alphas_cumprod", torch.log(1 - alphas_cumprod)) self.register("sqrt_recip_alphas_cumprod", torch.rsqrt(alphas_cumprod)) self.register("sqrt_recipm1_alphas_cumprod", torch.sqrt(1 / alphas_cumprod - 1)) self.register("posterior_variance", posterior_variance) self.register("posterior_log_variance_clipped", torch.log(torch.cat((posterior_variance[1].view(1, 1), posterior_variance[1:].view(-1, 1)), 0)).view(-1)) self.register("posterior_mean_coef1", (betas * torch.sqrt(alphas_cumprod_prev) / (1 - alphas_cumprod))) self.register("posterior_mean_coef2", ((1 - alphas_cumprod_prev) * torch.sqrt(alphas) / (1 - alphas_cumprod)))
def off_policy_weight(eval_log_p, behavior_log_p, full_trajectory=False, clamp_max=5.0): """Compute off-policy weight. Parameters ---------- eval_log_p: torch.tensor. Evaluation log probabilities. behavior_log_p: torch.tensor. Behavior log probabilities. full_trajectory: bool, optional (default=False). Flag that indicates whether the off-policy weight is for a single step or for the full trajectory. clamp_max: float. Value to clamp max. Returns ------- weight: torch.Tensor. Importance sample weights of the trajectory. """ weight = torch.exp(eval_log_p - behavior_log_p) if full_trajectory: weight = torch.cumprod(weight, dim=-1) return weight.clamp_max(clamp_max)
def _allocation(self, usage): usage = _EPSILON + (1 - _EPSILON) * usage sorted_usage, indices = torch.topk(usage, memory_size, largest=False, sorted=True) #this is not quite right, check prod_sorted_usage = torch.cumprod(sorted_usage, axis=1) exclusive_prod_sorted_usage = torch.ones(1, memory_size) exclusive_prod_sorted_usage[0][1:] = prod_sorted_usage[0][:-1] ##end sorted_allocation = (1 - sorted_usage) * exclusive_prod_sorted_usage final_allocation = torch.zeros((1, memory_size)) #undo the initial permutation for i in range(memory_size): a_index = indices[0][i] final_allocation[0][a_index] = sorted_allocation[0][i] return final_allocation
def _apply_discount(self, rewards): cum_discount = torch.cumprod( self.hp.gamma * torch.ones(*rewards[0].size()), dim=1) / self.hp.gamma discounted_rewards = rewards * cum_discount return discounted_rewards
def flat_to_net(flat_net): test_net = SimpleMNISTCNN() index = 0 for p in test_net.parameters(): end_index = torch.cumprod(torch.tensor(p.shape), 0)[-1] + index p = flat_net[index:end_index].reshape(p.shape) index = end_index
def forward(self, output, dropoutl): # output: [seq_len x batch_size x nhidlast] latent = self.latent(output) # h latent = self.lockdrop(latent, dropoutl) # h after variational dropout [seq_len x batch_size x n_experts * ninp] logit = self.decoder(latent.view(-1, self.ninp)) # HW [seq_len * batch_size * n_experts x voc_size] a = self.reduce(output.view(-1, self.nhidlast)) + 1e-8 # [seq_len * batch_size x n_experts] b = torch.sum(a, 1) c = torch.cumsum(a, 1) d = torch.abs(b[:, None] - c) a = a.to('cpu') d = d.to('cpu') beta = torch.distributions.Beta(a, d) sample = beta.rsample() # [seq_len * batch_size x n_experts] sample = sample.to('cuda') rem = 1 - sample D = torch.diag(torch.ones(self.n_experts - 1, device=device), 1) rem = rem @ D rem[:, 0] = 1 remprod = torch.cumprod(rem, 1) pis = remprod * sample prob = nn.functional.softmax(logit.view(-1, self.ntoken), dim=1).view(-1, self.n_experts, self.ntoken) # exp(hw) / sum(exp(hw)) prob = (prob * pis.unsqueeze(2).expand_as(prob)).sum(1) # weighted sum # TODO maybe we can do this with logsoftmax return prob
def _composite(self, optic_d, s_samples, rays_d): # distances between each samples dists = s_samples[..., 1:] - s_samples[ ..., :-1] # (chunk_size, N_samples - 1) dists = torch.cat( [dists, torch.tensor([1e10]).expand(dists[..., :1].shape)], -1) # (chunk_size, N_samples) dists = dists * torch.norm(rays_d[..., None, :], dim=-1) # retrieve display colors and alphas for each samples by a transfer function rgbs, alphas = self._transfer(optic_d, dists) weights = alphas * torch.cumprod(torch.cat( [torch.ones( (alphas.shape[0], 1)), 1.0 - alphas + 1e-10], dim=-1)[:, :-1], dim=-1) # (chunk_size, N_samples) rgb_map = torch.sum(weights[..., None] * rgbs, dim=-2) # (chunk_size, 3) acc_map = torch.sum(weights, -1) # (chunk_size) if self.config.white_bkgd: rgb_map = rgb_map + (1.0 - acc_map[..., None]) return rgb_map, weights
def raw2outputs_warp(raw_p, z_vals, rays_d, raw_noise_std=0): dists = z_vals[..., 1:] - z_vals[..., :-1] dists = torch.cat( [dists, torch.Tensor([1e10]).expand(dists[..., :1].shape)], -1) # [N_rays, N_samples] dists = dists * torch.norm(rays_d[..., None, :], dim=-1) rgb = torch.sigmoid(raw_p[..., :3]) # [N_rays, N_samples, 3] noise = 0. if raw_noise_std > 0.: noise = torch.randn(raw_p[..., 3].shape) * raw_noise_std act_fn = F.relu opacity = act_fn(raw_p[..., 3] + noise) alpha = 1. - torch.exp(-opacity * dists) weights = alpha * torch.cumprod( torch.cat([torch.ones( (alpha.shape[0], 1)), 1. - alpha + 1e-10], -1), -1)[:, :-1] rgb_map = torch.sum(weights[..., None] * rgb, -2) # [N_rays, 3] depth_map = torch.sum(weights * z_vals, -1) return rgb_map, depth_map, weights #, alpha #alpha#, 1. - probs
def _discount_rewards(self, rewards): gammas = torch.ones_like(rewards) if len(rewards) > 1: gammas[1:] = torch.cumprod(torch.tensor( self.config['discount_rate']).repeat(len(rewards) - 1), dim=0) return gammas * rewards
def _build_mask_and_subsample( self, not_dones, valids, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]: t = not_dones.size(1) valid_mask = (not_dones > 0) & valids valid_mask = valid_mask.float() valid_mask_unfolded = self._build_unfolded( valid_mask[:, :-1].to(dtype=torch.bool), self.k ) time_subsample = torch.randperm( t - 1, device=valid_mask.device, dtype=torch.long )[0:self.time_subsample] forward_mask = ( torch.cumprod(valid_mask_unfolded.index_select(2, time_subsample), dim=0) .to(dtype=torch.bool) .flatten(1, 2) ) max_k = forward_mask.flatten(1).any(-1).nonzero().max().item() + 1 unroll_subsample = torch.randperm(max_k, dtype=torch.long)[0:self.forward_subsample] max_k = unroll_subsample.max().item() + 1 unroll_subsample = unroll_subsample.to(device=valid_mask.device) forward_mask = forward_mask.index_select(0, unroll_subsample) return forward_mask, unroll_subsample, time_subsample, max_k
def make_objective(batch_data, v_fn, tau=1.0, dice_lambda=1.0, use_dice=False, gamma_weighted=True): """Make objective for DiCE, Loaded DiCE, or classic surrogate loss""" (batch_states, batch_pi_taken, batch_a_taken, batch_r, batch_dones, batch_mask, ep_returns) = batch_data empty_mask = (1 - batch_mask[:,:-1]).type(torch.ByteTensor) batch_pi_taken[empty_mask] = 1.0 log_pi = torch.log(batch_pi_taken) batch_values = v_fn(batch_states).detach() advantages = gae(batch_r, batch_values, batch_mask, tau=tau, gamma_weighted=gamma_weighted) if use_dice == "old": log_pi_cumsum = torch.cumsum(log_pi, 1) deps = magic_box(log_pi_cumsum) batch_r[:,-1] = batch_r[:,-1] + batch_values[:,-1]*gamma if gamma_weighted: gamma_weights = torch.cumprod(torch.ones_like(advantages) * gamma, 1) / gamma batch_r = batch_r * gamma_weights obj = (batch_r * deps).sum(1).mean() elif use_dice == "loaded": weighted_cumsum = torch.zeros_like(log_pi) weighted_cumsum[:,0] = log_pi[:,0] for t in range(1, log_pi.size(1)): weighted_cumsum[:,t] = dice_lambda * weighted_cumsum[:,t-1] + log_pi[:,t] deps_exclusive = weighted_cumsum - log_pi full_deps = magic_box(weighted_cumsum) - magic_box(deps_exclusive) obj = (advantages * full_deps).sum(1).mean() else: obj = (advantages * log_pi).sum(1).mean() return obj
def allocation_weighting(self): ''' Sorts the memory by usages first. Then perform calculation depending on the sort order. The alloation_weighting of the third least used memory is calculated as follows: Find the least used and second least used. Multiply their usages. Multiply the product with (1-usage of the third least), return. Do not confuse the sort order and the memory's natural location. Verify backprop. :param usage_vector: u_t, (N), [0,1] :return: allocation_wighting: a_t, (N), simplex bound ''' # not the last usage, since we will update usage before this sorted, indices = self.last_usage_vector.sort(dim=1) cum_prod = torch.cumprod(sorted, 1) # notice the index on the product cum_prod = torch.cat( [Variable(torch.ones(self.bs, 1).cuda()), cum_prod], 1)[:, :-1] sorted_inv = 1 - sorted allocation_weighting = sorted_inv * cum_prod # to shuffle back in place ret = torch.gather(allocation_weighting, 1, indices) if debug: if (ret != ret).any(): raise ValueError("NA found in allocation weighting") return ret
def dice_objective(self): self_logprobs = torch.stack(self.self_logprobs, dim=1) other_logprobs = torch.stack(self.other_logprobs, dim=1) values = torch.stack(self.values, dim=1) rewards = torch.stack(self.rewards, dim=1) # apply discount: cum_discount = torch.cumprod(hp.gamma * torch.ones(*rewards.size()), dim=1) / hp.gamma discounted_rewards = rewards * cum_discount discounted_values = values * cum_discount # stochastics nodes involved in rewards dependencies: dependencies = torch.cumsum(self_logprobs + other_logprobs, dim=1) # logprob of each stochastic nodes: stochastic_nodes = self_logprobs + other_logprobs # dice objective: dice_objective = torch.mean( torch.sum(magic_box(dependencies) * discounted_rewards, dim=1)) if hp.use_baseline: # variance_reduction: baseline_term = torch.mean( torch.sum( (1 - magic_box(stochastic_nodes)) * discounted_values, dim=1)) dice_objective = dice_objective + baseline_term return -dice_objective # want to minimize -objective
def forward(self, **kwargs): h = kwargs['x_path'] A, h = self.attention_net(h) A = torch.transpose(A, 1, 0) if 'attention_only' in kwargs.keys(): if kwargs['attention_only']: return A A_raw = A A = F.softmax(A, dim=1) M = torch.mm(A, h) logits = self.classifier(M) Y_hat = torch.topk(logits, 1, dim=1)[1] hazards = torch.sigmoid(logits) # hazards = F.softmax(logits, dim=1) S = torch.cumprod(1 - hazards, dim=1) # S = 1 - torch.cumsum(hazards, dim=1) results_dict = {} if 'return_features' in kwargs.keys(): if kwargs['return_features']: results_dict.update({'features': M}) return hazards, S, Y_hat, A_raw, results_dict
def jet_cross_entropy_loss(prediction: Tensor, target_data: Tensor, target_mask: Tensor, gamma: float) -> Tensor: batch_size = prediction.shape[0] prediction_shape = prediction.shape[1:] # Remove missing jets target_data = target_data.clamp(0, None) # Find the unravelling shape required to flatten the target indices ravel_sizes = torch.tensor(prediction_shape).flip(0) ravel_sizes = torch.cumprod(ravel_sizes, 0) ravel_sizes = ravel_sizes // ravel_sizes[0] ravel_sizes = ravel_sizes.flip(0).unsqueeze(0) ravel_sizes = ravel_sizes.to(target_data.device) # Flatten the target and predicted data to be one dimensional ravel_target = (target_data * ravel_sizes).sum(1) ravel_prediction = prediction.reshape(batch_size, -1).contiguous() log_probability = ravel_prediction.gather(-1, ravel_target.view(-1, 1)).squeeze() focal_scale = (1 - torch.exp(log_probability))**gamma return -log_probability * focal_scale * target_mask
def render(self, volume, axis=2): padding = 1 volume = torch.nn.functional.pad( volume, (padding, padding, padding, padding, padding, padding)) density = volume[:, [3]] signal = volume[:, :3] bs = density.shape[0] sample_coords = self.sample_coords.expand(bs, self.sample_size_z, self.param.data.cube_len, self.param.data.cube_len, 3) density = density * self.density_factor density = density / self.sample_size_z density = torch.nn.functional.grid_sample(density, sample_coords) transmission = torch.cumprod(1.0 - density, dim=axis) weight = density * transmission weight_sum = torch.sum(weight, dim=axis) signal = torch.nn.functional.grid_sample(signal, sample_coords) rendering = torch.sum(weight * signal, dim=axis) rendering = rendering / (weight_sum + 1e-8) alpha = 1.0 - torch.prod(1 - density, dim=axis) rendering = rendering * alpha rendering = torch.cat([rendering, alpha], dim=1) return rendering
def retrace(q_values, rewards, importance_weights=None, discount=0.99, l=0.75): """ Retrace estimate (Munos et al., 2016). Args: q_values (torch.Tensor): the Q-value estimates at each time step [time_steps+1, batch_size, 1] rewards (torch.Tensor): the rewards at each time step [time_steps, batch_size, 1] importance_weights (torch.Tensor): the importance weights at each time step [time_steps, batch_size, 1] discount (float): the temporal discount factor l (float): the lambda weighting factor """ if q_values.shape[0] == 1 or l == -1: # degenerate case return q_values if importance_weights is None: # On-policy importance_weights = torch.ones_like(q_values) deltas = rewards + discount * q_values[1:] - q_values[:-1] importance_weights = l * torch.clamp(importance_weights, 0, 1)[:-2] importance_weights = torch.cat([torch.ones_like(q_values[:1]), importance_weights], 0) discounts = torch.cat([(discount*torch.ones_like(q_values[:1]))**i for i in range(q_values.shape[0])], 0) q_estimates = q_values[:1] + torch.sum(discounts[:-1] * torch.cumprod(importance_weights, 0) * deltas, 0, keepdim=True) return q_estimates
def forward(self, x): x0 = self.conv.forward(x.float()) x = self.pool_mil(x0) x = x.squeeze(2).squeeze(2) x1 = torch.add(torch.mul(x0.view(x.size(0), 1000, -1), -1), 1) cumprod = torch.cumprod(x1, 2) out = torch.max(x, torch.add(torch.mul(cumprod[:, :, -1], -1), 1)) #out = F.softmax(out) return out
def forward(self, img, att_size=14): x0 = self.conv(img) x = self.pool_mil(x0) x = x.squeeze(2).squeeze(2) x = self.l1(x) x1 = torch.add(torch.mul(x.view(x.size(0), 1000, -1), -1), 1) cumprod = torch.cumprod(x1, 2) out = torch.max(x, torch.add(torch.mul(cumprod[:, :, -1], -1), 1)) return out
def forward(ctx, input, dim): ctx.dim = dim ctx.save_for_backward(input) return torch.cumprod(input, dim=ctx.dim)
def backward(ctx, grad_output): ''' There are two algorithms to do this. The first one is very efficient, but works only when there are no nonzero elements in the input. The second one is much more complex, but it doesn't assume anything on the input. The main downside is that it takes time O(n^2), where n = input.size(self.dim) (i.e. the length of the cumulative product). This is in contrast to the forward pass and the efficient algorithm, which are both O(n). The second algorithm is a simple application of the chain rule. If x is an n-dimensional vector, and y = cumprod(x), and F is the final cost, then dF / dx_k = sum_j (dF / dy_j) * (dy_j / dx_k) (1) The term dF / dy_j is just grad_output[j] (assuming again everything is one-dimensional). The term (dy_j / dx_k) is easilly seen to be if j >= k dy_j / dx_k = prod_{1 <= i <= j, i != k} x_i else: dy_j / dx_k = 0 Note that the indicator (j>=k) can be taken out by replacing the sum in (1) with a sum from j = k to n. Thus, df / dx_k = sum_{k <= j <= n} grad_output[j] * (dy_j / dx_k) with dy_j / dx_k = prod_{1 <= i <= j, i != k} x_i (2) Note that this last term is just the cumulative product with k omitted. Thus, if x_k (the input) is nonzero, we can just express this as dy_j / dx_k = (prod_{1 <= i <= j} x_i) / x_k = y_j / x_k So therefore, df / dx_k = sum_{k <= j <= n} grad_output[j] * y_j / x_k so grad_output = sum_scan_exclusiv(grad_output * output) / input If the input is nonzero, we need to calculate the dy_j / dx_k by using the formula (2), called in the code omitted_products. The way the code calculates it is simply by noting that prod_{1 <= i <= j, i != k} x_i = (prod_{1 <= i <= k} x_i) * (prod_{k + 1 <= i <= j} x_i) the first term is calculated as prods_until_k, which since doesn't depend in j is easy to vectorize. The second term (indexed by j) is the cumulative product of x_{k+1}, x_{k+2}, ..., x_n, and it's named in the code prods_from_k_pkus_1, and it's calculated as a cumprod. In order to vectorize this properly, we need to add to omitted_products the dimensions where k > j, and therefore dy_j / dx_k = 0, which is done right after the assert. ''' input, = ctx.saved_variables dim_size = input.size(ctx.dim) if dim_size == 1: return grad_output, None # Simple case with nonzero elements in the input if (input != 0).data.all(): output = torch.cumprod(input, dim=ctx.dim) return sum_scan_exclusive(output * grad_output, dim=ctx.dim) / input, None positive_dim = ctx.dim if ctx.dim >= 0 else input.dim() + ctx.dim dim_padding = (slice(None, None),) * (positive_dim) ones_size = list(input.size()) ones_size[ctx.dim] = 1 ones = Variable(input.data.new([1]).expand(ones_size)) grad_input = Variable(grad_output.data.new(input.size()).zero_()) for k in range(dim_size): if k == 0: prods_from_k_plus_1 = torch.cumprod( input[dim_padding + (slice(k + 1, None),)], dim=ctx.dim ) omitted_products = torch.cat( (ones, prods_from_k_plus_1), dim=ctx.dim ) elif k == dim_size - 1: prods_until_k = torch.prod( input[dim_padding + (slice(None, k),)], dim=ctx.dim, keepdim=True ) omitted_products = prods_until_k else: prods_until_k = torch.prod( input[dim_padding + (slice(None, k),)], dim=ctx.dim, keepdim=True ) prods_from_k_plus_1 = torch.cumprod( input[dim_padding + (slice(k + 1, None),)], dim=ctx.dim ) omitted_products = prods_until_k.expand_as( prods_from_k_plus_1) * prods_from_k_plus_1 omitted_products = torch.cat( (prods_until_k, omitted_products), ctx.dim) # At this point omitted_products is the same size # as input, except on the dimension dim where it's # dim_size - k assert omitted_products.size(ctx.dim) == dim_size - k # should we implement copy_ or _set_item in variable? index = tuple(slice(None, None) for _ in range(positive_dim)) + (k,) grad_input[index] = torch.sum( grad_output[dim_padding + (slice(k, None),)] * omitted_products, dim=ctx.dim) return grad_input, None