def check(self, h, w, cs, rs, pa, rtp, dim): a = jt.random([h, w]) a.data with jt.log_capture_scope( log_v=0, log_vprefix="tuner_manager=100", # this value is used for force compile compile_options={"test_reduce_tuner": 1}) as logs: amean = jt.mean(a, dims=[dim], keepdims=1) a2mean = jt.mean(a * a, dims=[dim], keepdims=1) norm_aa = (a - amean.broadcast_var(a)) / ( jt.sqrt(a2mean - amean * amean).broadcast_var(a)) norm_aa.data logs = find_log_with_re( logs, "Run tuner reduce: confidence\\((20)\\) candidates\\((.*)\\)$") assert len(logs) == 1, logs assert logs[0][0] == "20", "confidence of reorder should be 20" candidates = simple_parser(logs[0][1]) assert candidates == { "order0": [ 0, ], "order1": [ 1, ], "order2": [ 0, ], "split1": [ 2048, ], }
def execute(self, x): dims = [0] + list(range(2, x.ndim)) ####### centering calibration begin ####### $ x += self.center_weight * self.stas(x) ####### centering calibration end ####### $ if self.is_train: xmean = jt.mean(x, dims=dims) x2mean = jt.mean(x * x, dims=dims) if self.sync and jt.in_mpi: xmean = xmean.mpi_all_reduce("mean") x2mean = x2mean.mpi_all_reduce("mean") xvar = (x2mean - xmean * xmean).maximum(0.0) w = 1.0 / jt.sqrt(xvar + self.eps) b = -xmean * w norm_x = x * w.broadcast(x, dims) + b.broadcast(x, dims) self.running_mean.update(self.running_mean + (xmean.reshape( (-1, )) - self.running_mean) * self.momentum) self.running_var.update(self.running_var + (xvar.reshape((-1, )) - self.running_var) * self.momentum) else: w = 1.0 / jt.sqrt(self.running_var + self.eps) b = -self.running_mean * w norm_x = x * w.broadcast(x, dims) + b.broadcast(x, dims) ####### scaling calibration begin ####### $ scale_factor = jt.sigmoid(self.scale_weight * self.stas(norm_x) + self.scale_bias) ####### scaling calibration end ####### $ return self.weight * scale_factor * norm_x + self.bias
def execute(self, x): if len(x.shape) == 3: dims = [0, 2] else: dims = [0] if self.is_train: xmean = jt.mean(x, dims=dims, keepdims=1) x2mean = jt.mean(x * x, dims=dims, keepdims=1) if self.sync and jt.in_mpi: xmean = xmean.mpi_all_reduce("mean") x2mean = x2mean.mpi_all_reduce("mean") xvar = x2mean - xmean * xmean norm_x = (x - xmean) / jt.sqrt(xvar + self.eps) self.running_mean.update(self.running_mean + (xmean.sum(dims) - self.running_mean) * self.momentum) self.running_var.update(self.running_var + (xvar.sum(dims) - self.running_var) * self.momentum) else: running_mean = self.running_mean.broadcast(x, dims) running_var = self.running_var.broadcast(x, dims) norm_x = (x - running_mean) / jt.sqrt(running_var + self.eps) if not self.affine: return norm_x w = self.weight.broadcast(x, dims) b = self.bias.broadcast(x, dims) return norm_x * w + b
def dis_loss(self, real_samps, fake_samps, height, alpha): r_preds = self.dis(real_samps, height, alpha) f_preds = self.dis(fake_samps, height, alpha) # loss = (torch.mean(nn.ReLU()(1 - r_preds)) + # torch.mean(nn.ReLU()(1 + f_preds))) loss = (jt.mean(nn.ReLU()(1 - r_preds)) + jt.mean(nn.ReLU()(1 + f_preds))) return loss
def batch_norm(x): xmean = jt.mean(x, dims=[0, 2, 3], keepdims=1) x2mean = jt.mean(x * x, dims=[0, 2, 3], keepdims=1) norm_x = (x - xmean.broadcast_var(x)) / ( jt.sqrt(x2mean - xmean * xmean + jt.float32(1e-5)).broadcast_var(x)) w = jt.make_var([x.shape[1]], init=get_init_var) b = jt.make_var([x.shape[1]], init=get_init_var) w = w.broadcast([1, w.shape[0], 1, 1], [0, 2, 3]) b = b.broadcast([1, b.shape[0], 1, 1], [0, 2, 3]) return norm_x * w + b
def execute(self, x): dims = [-i for i in range(len(self.normalized_shape), 0, -1)] mean = jt.mean(x, dims=dims, keepdims=1) numerator = x - mean variance = jt.mean(numerator.sqr(), dims=dims, keepdims=1) denominator = jt.sqrt(variance + self.eps) norm_x = numerator / denominator if self.elementwise_affine: norm_x = norm_x * self.weight + self.bias return norm_x
def execute(self, x): xmean = jt.mean(x, dims=[2, 3], keepdims=1) x2mean = jt.mean(x * x, dims=[2, 3], keepdims=1) if self.sync and jt.in_mpi: xmean = xmean.mpi_all_reduce("mean") x2mean = x2mean.mpi_all_reduce("mean") xvar = jt.maximum(x2mean - xmean * xmean, 0) norm_x = (x - xmean) / jt.sqrt(xvar + self.eps) w = self.weight.broadcast(x, [0, 2, 3]) b = self.bias.broadcast(x, [0, 2, 3]) return norm_x * w + b
def execute(self, x): N, C, H, W = x.shape assert C == self.num_channels assert C % self.num_groups == 0 x = x.reshape((N, self.num_groups, int(C / self.num_groups), H * W)) xmean = jt.mean(x, dims=[2, 3], keepdims=1) x2mean = jt.mean(x * x, dims=[2, 3], keepdims=1) xvar = jt.maximum(x2mean - xmean * xmean, 0) norm_x = (x - xmean) / jt.sqrt(xvar + self.eps) w = self.weight.reshape((1, self.num_groups, C // self.num_groups, 1)) b = self.bias.reshape((1, self.num_groups, C // self.num_groups, 1)) return (norm_x * w + b).reshape((N, C, H, W))
def execute(self, x): if self.is_train: xmean = jt.mean(x, dims=[0,2,3], keepdims=1) x2mean = jt.mean(x*x, dims=[0,2,3], keepdims=1) xvar = x2mean-xmean*xmean norm_x = (x-xmean)/jt.sqrt(xvar+self.eps) self.running_mean += (xmean.sum([0,2,3])-self.running_mean)*self.momentum self.running_var += (xvar.sum([0,2,3])-self.running_var)*self.momentum else: running_mean = self.running_mean.broadcast(x, [0,2,3]) running_var = self.running_var.broadcast(x, [0,2,3]) norm_x = (x-running_mean)/jt.sqrt(running_var+self.eps) w = self.weight.broadcast(x, [0,2,3]) b = self.bias.broadcast(x, [0,2,3]) return norm_x * w + b
def dis_loss(self, real_samps, fake_samps, height, alpha, r1_gamma=10.0): # Obtain predictions r_preds = self.dis(real_samps, height, alpha) f_preds = self.dis(fake_samps, height, alpha) # loss = torch.mean(nn.Softplus()(f_preds)) + torch.mean(nn.Softplus()(-r_preds)) loss = jt.mean(nn.Softplus()(f_preds)) + jt.mean( nn.Softplus()(-r_preds)) if r1_gamma != 0.0: r1_penalty = self.R1Penalty(real_samps.detach(), height, alpha) * (r1_gamma * 0.5) loss += r1_penalty return loss
def chamfer_loss(pc1, pc2, reduction='mean', sqrt=True): ''' return the chamfer loss from pc1 to pc2. Parameters: =========== pc1: [B, N, xyz] pc2: [B, N, xyz] reduction: 'mean', 'sum', or None ''' batch_size_1, n_samples_pc1, _ = pc1.shape batch_size_2, n_samples_pc2, _ = pc2.shape assert batch_size_1 == batch_size_2 batch_size = batch_size_1 idx = jt.code([batch_size, n_samples_pc1], 'int32', [pc1, pc2], cpu_src=cpu_src, cuda_src=cuda_src) nearest_pts = select_vertices(pc2, idx) if sqrt: chamfer_distance = (((pc1 - nearest_pts)**2).sum(dim=-1)).sqrt() else: chamfer_distance = (((pc1 - nearest_pts)**2).sum(dim=-1)) if reduction is None: return chamfer_distance elif reduction == 'sum': return jt.sum(chamfer_distance) elif reduction == 'mean': return jt.mean(chamfer_distance)
def execute(self, x): avg_out = jt.mean(x, dim=1, keepdims=1) # 压缩通道 max_out = jt.max(x, dim=1, keepdims=1) # 压缩通道 x = jt.contrib.concat([avg_out, max_out], dim=1) # [b, 1, h, w] x = self.conv1(x) y = self.sigmoid(x) return y
def execute(self, x, cls_label): batch_size, _, N = x.size() x = self.relu(self.bn1(self.conv1(x))) # B, D, N x = self.relu(self.bn2(self.conv2(x))) x1 = self.sa1(x) x2 = self.sa2(x1) x3 = self.sa3(x2) x4 = self.sa4(x3) x = concat((x1, x2, x3, x4), dim=1) x = self.conv_fuse(x) x_max = jt.max(x, 2) x_avg = jt.mean(x, 2) x_max_feature = x_max.view(batch_size, -1).unsqueeze(-1).repeat(1, 1, N) x_avg_feature = x_avg.view(batch_size, -1).unsqueeze(-1).repeat(1, 1, N) cls_label_one_hot = cls_label.view(batch_size, 16, 1) cls_label_feature = self.label_conv(cls_label_one_hot).repeat(1, 1, N) x_global_feature = concat( (x_max_feature, x_avg_feature, cls_label_feature), 1) # 1024 + 64 x = concat((x, x_global_feature), 1) # 1024 * 3 + 64 x = self.relu(self.bns1(self.convs1(x))) x = self.dp1(x) x = self.relu(self.bns2(self.convs2(x))) x = self.convs3(x) return x
def check(self, h, w, cs, rs, pa, rtp, dim): a = jt.random([h, w]) a.sync() with jt.log_capture_scope( log_v=0, log_vprefix="tuner_manager=100", # this value is used for force compile compile_options={"test_new_fused_op": 1}) as logs: amean = jt.mean(a, dims=[dim], keepdims=1) a2mean = jt.mean(a * a, dims=[dim], keepdims=1) norm_aa = (a - amean.broadcast_var(a)) / ( jt.sqrt(a2mean - amean * amean).broadcast_var(a)) norm_aa.sync() logs = find_log_with_re( logs, "Run tuner reduce: confidence\\((.*)\\) candidates\\((.*)\\)$") assert len(logs) == 3, logs
def execute(self, x): N = x.shape[0] C = self.num_channels output_shape = (N, -1) # TODO: 3d group norm if x.ndim == 4: output_shape = x.shape assert C % self.num_groups == 0 x = x.reshape((N, self.num_groups, int(C / self.num_groups), -1)) xmean = jt.mean(x, dims=[2, 3], keepdims=1) x2mean = jt.mean(x * x, dims=[2, 3], keepdims=1) xvar = jt.maximum(x2mean - xmean * xmean, 0) norm_x = (x - xmean) / jt.sqrt(xvar + self.eps) if not self.affine: return norm_x.reshape(output_shape) w = self.weight.reshape((1, self.num_groups, C // self.num_groups, 1)) b = self.bias.reshape((1, self.num_groups, C // self.num_groups, 1)) return (norm_x * w + b).reshape(output_shape)
def gen_loss(self, real_samps, fake_samps, height, alpha): # Obtain predictions r_preds = self.dis(real_samps, height, alpha) f_preds = self.dis(fake_samps, height, alpha) # difference between real and fake: # r_f_diff = r_preds - torch.mean(f_preds) r_f_diff = r_preds - jt.mean(f_preds) # difference between fake and real samples # f_r_diff = f_preds - torch.mean(r_preds) f_r_diff = f_preds - jt.mean(r_preds) # return the loss # return (torch.mean(nn.ReLU()(1 + r_f_diff)) # + torch.mean(nn.ReLU()(1 - f_r_diff))) return (jt.mean(nn.ReLU()(1 + r_f_diff)) + jt.mean(nn.ReLU()(1 - f_r_diff)))
def chamfer_loss(pc1, pc2, reduction='mean', dims='BNC', bidirectional=False): ''' return the chamfer loss from pc1 to pc2. :param pc1: input point cloud :type pc1: jittor array :param pc2: input point cloud :type pc2: jittor array :param reduction: reduction method in batches, can be 'mean', 'sum', or None. Default: 'mean'. :type reduction: str, optional :param dims: a string that represents each dimension, can be '[BNC]' ([batch, number of points, xyz]), or '[BCN]' ([batch, xyz, number of points]). Default: 'BNC'. :type dims: str, optional Example: >>> import jittor as jt >>> from jittor.loss3d import chamfer_loss >>> jt.flags.use_cuda = True >>> pc1 = jt.rand([10, 100, 3], dtype=jt.float32) >>> pc2 = jt.rand([10, 100, 3], dtype=jt.float32) >>> cf = chamfer_loss(pc1, pc2, dims='BNC', bidirectional=True) >>> print('chamfer loss =', cf.item()) ''' if bidirectional: return chamfer_loss(pc1, pc2, reduction, dims) + chamfer_loss( pc2, pc1, reduction, dims) assert dims in ['BNC', 'BCN'] if dims == 'BCN': pc1, pc2 = pc1.permute(0, 2, 1), pc2.permute(0, 2, 1) batch_size_1, N, _ = pc1.shape batch_size_2, M, _ = pc2.shape assert batch_size_1 == batch_size_2 batch_size = batch_size_1 idx = jt.code([batch_size, N], 'int32', [pc1, pc2], cpu_src=cpu_src, cuda_src=cuda_src) nearest_pts = pc2.reindex([batch_size, idx.shape[1], 3], ['i0', '@e0(i0, i1)', 'i2'], extras=[idx]) chamfer_distance = (((pc1 - nearest_pts)**2).sum(dim=-1)).sqrt() if reduction is None: return chamfer_distance elif reduction == 'sum': return jt.sum(chamfer_distance) elif reduction == 'mean': return jt.mean(chamfer_distance)
def execute(self, x): try: avg_out = jt.mean(x, dim=1, keepdims=True) max_out = jt.max(x, dim=1, keepdims=True) scale = jt.contrib.concat([avg_out, max_out], dim=1) scale = self.conv(scale) out = x * self.sigmoid(scale) except Exception as e: print(e) out = x return out
def execute(self, x): if self.is_train: xmean = jt.mean(x, dims=[0,2,3], keepdims=1) x2mean = jt.mean(x*x, dims=[0,2,3], keepdims=1) if self.sync and jt.in_mpi: xmean = xmean.mpi_all_reduce("mean") x2mean = x2mean.mpi_all_reduce("mean") xvar = x2mean-xmean*xmean norm_x = (x-xmean)/jt.sqrt(xvar+self.eps) self.running_mean.update(self.running_mean + (xmean.reshape((-1,)) - self.running_mean) * self.momentum) self.running_var.update(self.running_var + (xvar.reshape((-1,))-self.running_var)*self.momentum) else: running_mean = self.running_mean.broadcast(x, [0,2,3]) running_var = self.running_var.broadcast(x, [0,2,3]) norm_x = (x-running_mean)/jt.sqrt(running_var+self.eps) w = self.weight.broadcast(x, [0,2,3]) b = self.bias.broadcast(x, [0,2,3]) return norm_x * w + b
def batch_norm(x, is_train, eps=1e-5, momentum=0.1): w = jt.make_var([x.shape[1]], init=lambda *a: init.constant(*a, 1.0)) b = jt.make_var([x.shape[1]], init=lambda *a: init.constant(*a, 0.0)) running_mean = jt.make_var([x.shape[1]], init=lambda *a: init.constant(*a, 0.0)) running_var = jt.make_var([x.shape[1]], init=lambda *a: init.constant(*a, 1.0)) w = w.broadcast(x, [0,2,3]) b = b.broadcast(x, [0,2,3]) if is_train: xmean = jt.mean(x, dims=[0,2,3], keepdims=1) x2mean = jt.mean(x*x, dims=[0,2,3], keepdims=1) xvar = x2mean-xmean*xmean norm_x = (x-xmean)/jt.sqrt(xvar+eps) running_mean += (xmean.sum([0,2,3])-running_mean)*momentum running_var += (xvar.sum([0,2,3])-running_var)*momentum else: running_mean = running_mean.broadcast(x, [0,2,3]) running_var = running_var.broadcast(x, [0,2,3]) norm_x = (x-running_mean)/jt.sqrt(running_var+eps) return norm_x * w + b
def execute(self, x, normal=None): if normal is None: x = (x, x) else : x = (x, normal) x = self.pcnn1(x) x = self.pcnn2(x)[1] # grab features x = x.permute(0, 2, 1) # b, dim, n logits = self.fcn(x) logits = jt.mean(logits, dim=2) return logits
def execute(self, trans_points, cp, voxel, gridSize, weight=1): if len(trans_points.shape) == 4: trans_points = trans_points.squeeze(dim=-1) nb = pointClosestCellIndex(trans_points) idx = jt.matmul( nb, jt.transform.to_tensor(jt.array([(gridSize**2), gridSize, 1]))) mask = (1 - voxel.view((-1), (gridSize**3)).gather(1, idx)) idx = idx.unsqueeze(2) idx = idx.repeat(1, 1, 3) mask = mask.unsqueeze(2).repeat(1, 1, 3) closest_points = cp.gather(1, idx) self.constant = weight distance = (trans_points - closest_points) distance = (distance * mask) # self.save_for_backward(distance) self.saved_tensors = distance return (jt.mean(jt.sum(jt.sum(jt.pow(distance, 2), 2), 1)) * weight)
def adjust_contrast_tensor(img, contrast_factor): """Adjust contrast of an RGB image. Args: img (Tensor): Image to be adjusted. contrast_factor (float): How much to adjust the contrast. Can be any non negative number. 0 gives a solid gray image, 1 gives the original image while 2 increases the contrast by a factor of 2. Returns: Tensor: Contrast adjusted image. """ if contrast_factor < 0: raise ValueError('contrast_factor ({}) is not non-negative.'.format(contrast_factor)) if not _is_tensor_a_jittor_image(img): raise TypeError('tensor is not a jittor image.') gray = rgb_to_grayscale(img) gray.dtype = 'float' mean = jt.mean(gray) return _blend(img, mean, contrast_factor)
def gen_loss(self, _, fake_samps, height, alpha): f_preds = self.dis(fake_samps, height, alpha) # print(f_preds.is_stop_grad()) # return torch.mean(nn.Softplus()(-f_preds)) return jt.mean(nn.Softplus()(-f_preds))
def execute(self, input): return input / jt.sqrt(jt.mean(input ** 2, dim=1, keepdims=True) + 1e-8)
def gen_loss(self, _, fake_samps, height, alpha): # return -torch.mean(self.dis(fake_samps, height, alpha)) return -jt.mean(self.dis(fake_samps, height, alpha))
latent_dim=latent_dim, n_c=n_c) gen_imgs = generator(zn, zc) D_gen = discriminator(gen_imgs) D_real = discriminator(real_imgs) # ----------------- # Train Generator # ----------------- if ((i % n_skip_iter) == 0): (enc_gen_zn, enc_gen_zc, enc_gen_zc_logits) = encoder(gen_imgs) zn_loss = mse_loss(enc_gen_zn, zn) zc_loss = xe_loss(enc_gen_zc_logits, zc_idx) if wass_metric: ge_loss = ((jt.mean(D_gen) + (betan * zn_loss)) + (betac * zc_loss)) else: valid = jt.ones([gen_imgs.shape[0], 1]).stop_grad() v_loss = bce_loss(D_gen, valid) ge_loss = ((v_loss + (betan * zn_loss)) + (betac * zc_loss)) optimizer_GE.step(ge_loss) # --------------------- # Train Discriminator # --------------------- if wass_metric: grad_penalty = calc_gradient_penalty(discriminator, real_imgs, gen_imgs) d_loss = ((jt.mean(D_real) - jt.mean(D_gen)) + grad_penalty)
total_time = 0. cnt = 0 # ---------- # Training # ---------- for epoch in range(opt.n_epochs): for (i, (real_imgs, _)) in enumerate(dataloader): # ----------------- # Train Discriminator # ----------------- z = jt.array(np.random.normal(0, 1, (real_imgs.shape[0], opt.latent_dim)).astype(np.float32)) fake_imgs = generator(z).detach() loss_D = ((- jt.mean(discriminator(real_imgs))) + jt.mean(discriminator(fake_imgs))) optimizer_D.step(loss_D) for p in discriminator.parameters(): clamp_(p, - opt.clip_value, opt.clip_value) # --------------------- # Train Generator # --------------------- if ((i % opt.n_critic) == 0): gen_imgs = generator(z) loss_G = (- jt.mean(discriminator(gen_imgs))) optimizer_G.step(loss_G) if warmup_times==-1: print(('[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f]' % (epoch, opt.n_epochs, (batches_done % len(dataloader)), len(dataloader), loss_D.numpy()[0], loss_G.numpy()[0])))
(fake_B_gray_line2 >= 0)].sum()) / bs * opt.lambda_chamfer2 # Local loss real_B_locals = [ real_B_eyel, real_B_eyer, real_B_nose, real_B_mouth, real_B_hair, real_B_bg ] loss_G_local = 0 for j in range(6): loss_G_local += criterion_pixelwise( fake_B_locals[j], real_B_locals[j]) * opt.lambda_local # Line continuity loss fake_B_patches, conti_weights = get_patches(fake_B, maskface) outputs = regressor(fake_B_patches) loss_G_continuity = jt.mean( (1.0 - outputs) * conti_weights) * opt.lambda_continuity # Total loss loss_G = loss_GAN + loss_GAN_local + loss_pixel + ( loss_G_chamfer + loss_G_chamfer2) + loss_G_local + loss_G_continuity #pdb.set_trace() optimizer_G.step(loss_G) # --------------------- # Train Discriminator # --------------------- # Real loss pred_real = D_global(real_B, real_A) loss_real = criterion_GAN(pred_real, valid) loss_real_local = 0
def entropy(self): return -jt.sum(jt.mean(self.probs) * jt.log(self.probs))