def sample_pdf(bins, weights, N_samples, det=False): # Get pdf weights = weights + 1e-5 # prevent nans pdf = weights / jt.sum(weights, -1, keepdims=True) cdf = jt.cumsum(pdf, -1) cdf = jt.concat([jt.zeros_like(cdf[..., :1]), cdf], -1) # (batch, len(bins)) # Take uniform samples if det: u = jt.linspace(0., 1., steps=N_samples) u = u.expand(list(cdf.shape[:-1]) + [N_samples]) else: u = jt.random(list(cdf.shape[:-1]) + [N_samples]) # Invert CDF inds = jt.searchsorted(cdf, u, right=True) below = jt.maximum(jt.zeros_like(inds - 1), inds - 1) above = jt.minimum((cdf.shape[-1] - 1) * jt.ones_like(inds), inds) inds_g = jt.stack([below, above], -1) # (batch, N_samples, 2) matched_shape = [inds_g.shape[0], inds_g.shape[1], cdf.shape[-1]] cdf_g = jt.gather(cdf.unsqueeze(1).expand(matched_shape), 2, inds_g) bins_g = jt.gather(bins.unsqueeze(1).expand(matched_shape), 2, inds_g) denom = (cdf_g[..., 1] - cdf_g[..., 0]) denom[denom < 1e-5] = 1.0 t = (u - cdf_g[..., 0]) / denom samples = bins_g[..., 0] + t * (bins_g[..., 1] - bins_g[..., 0]) return samples
def __call__(self, matched_idxs): """ Arguments: matched idxs: list of tensors containing -1, 0 or positive values. Each tensor corresponds to a specific image. -1 values are ignored, 0 are considered as negatives and > 0 as positives. Returns: pos_idx (list[tensor]) neg_idx (list[tensor]) Returns two lists of binary masks for each image. The first list contains the positive elements that were selected, and the second list the negative example. """ pos_idx = [] neg_idx = [] for matched_idxs_per_image in matched_idxs: positive = jt.nonzero(matched_idxs_per_image >= 1).squeeze(1) negative = jt.nonzero(matched_idxs_per_image == 0).squeeze(1) num_pos = int(self.batch_size_per_image * self.positive_fraction) # protect against not enough positive examples num_pos = min(positive.numel(), num_pos) num_neg = self.batch_size_per_image - num_pos # protect against not enough negative examples num_neg = min(negative.numel(), num_neg) # randomly select positive and negative examples perm1 = jt.randperm(positive.numel())[:num_pos] perm2 = jt.randperm(negative.numel())[:num_neg] pos_idx_per_image = positive[perm1] neg_idx_per_image = negative[perm2] # create binary mask from indices pos_idx_per_image_mask = jt.zeros_like( matched_idxs_per_image).bool() neg_idx_per_image_mask = jt.zeros_like( matched_idxs_per_image).bool() pos_idx_per_image_mask[pos_idx_per_image] = 1 neg_idx_per_image_mask[neg_idx_per_image] = 1 pos_idx.append(pos_idx_per_image_mask) neg_idx.append(neg_idx_per_image_mask) return pos_idx, neg_idx
def semantic_segmentation_loss(self, segment_data, mask_t, class_t, interpolation_mode='bilinear'): # Note num_classes here is without the background class so cfg.num_classes-1 batch_size, num_classes, mask_h, mask_w = segment_data.shape loss_s = 0 for idx in range(batch_size): cur_segment = segment_data[idx] cur_class_t = class_t[idx] with jt.no_grad(): downsampled_masks = nn.interpolate( mask_t[idx].unsqueeze(0), (mask_h, mask_w), mode=interpolation_mode, align_corners=False).squeeze(0) downsampled_masks = (downsampled_masks > 0.5).float() # Construct Semantic Segmentation segment_t = jt.zeros_like(cur_segment) segment_t.stop_grad() for obj_idx in range(downsampled_masks.shape[0]): segment_t[cur_class_t[obj_idx]] = jt.maximum( segment_t[cur_class_t[obj_idx]], downsampled_masks[obj_idx]) loss_s += nn.BCEWithLogitsLoss(size_average=False)(cur_segment, segment_t) return loss_s / mask_h / mask_w * cfg.semantic_segmentation_alpha
def pre_step(self, loss): """ something should be done before step, such as calc gradients, mpi sync, and so on. Example:: class MyOptimizer(Optimizer): def step(self, loss): self.post_step(loss) ... """ # clean prev grads params = [] params_has_grad = [] for pg in self.param_groups: for p in pg['params']: params.append(p) if not p.is_stop_grad(): params_has_grad.append(p) # get gradient grads = jt.grad(loss, params_has_grad) # sync grads and model if in mpi if jt.in_mpi: dep = [] def add_dep(v): nonlocal dep v._add_dependency(dep) dep = [v] for g in grads: g.assign(g.mpi_all_reduce("mean")) add_dep(g._input(0)) if self.n_step % self.param_sync_iter == 0: for p in params: p.assign(p.mpi_broadcast()) add_dep(p) self.n_step += 1 # set up grads in param_groups pid = 0 for pg in self.param_groups: if "grads" not in pg: pg["grads"] = [ jt.zeros_like(p).stop_grad().stop_fuse() for p in pg['params'] ] pg_grads = pg["grads"] for i, p in enumerate(pg['params']): if not p.is_stop_grad(): # accumulate grad and stop grad of grad g = grads[pid].stop_grad() if not self.__zero_grad: g = g + pg_grads[i] pg_grads[i].update(g) pid += 1 self.__zero_grad = False
def test_twobilinear_lstm(self): x = jt.rand(5, 4, 10) rnn1 = nn.LSTM(10, 20, bidirectional=True) out1, _ = rnn1(x) rnn2 = nn.LSTM(40, 20, bidirectional=True) out2, _ = rnn2(out1) target = jt.zeros_like(out2) loss = nn.mse_loss(out2, target) from jittor import optim optimizer = optim.RMSprop(rnn1.parameters()) optimizer.step(loss)
def expand_boxes(boxes, scale): w_half = (boxes[:, 2] - boxes[:, 0]) * .5 h_half = (boxes[:, 3] - boxes[:, 1]) * .5 x_c = (boxes[:, 2] + boxes[:, 0]) * .5 y_c = (boxes[:, 3] + boxes[:, 1]) * .5 w_half *= scale h_half *= scale boxes_exp = jt.zeros_like(boxes) boxes_exp[:, 0] = x_c - w_half boxes_exp[:, 2] = x_c + w_half boxes_exp[:, 1] = y_c - h_half boxes_exp[:, 3] = y_c + h_half return boxes_exp
def reflect_coordinates(x, twice_low, twice_high): if twice_low == twice_high: return jt.zeros_like(x) m = twice_low / 2 span = (twice_high - twice_low) / 2 x = (x - m).abs() #`fmod` returns same sign as `in`, which is positive after the `fabs` above. extra = x.mod(span) flips = (x / span).floor() result1 = extra + m result2 = span - extra + m con = flips % 2 == 0 not_con = flips % 2 != 0 result1[not_con] = 0.0 result2[con] = 0.0 return result1 + result2
def voxelize_sub3(faces, voxels): bs = voxels.size(0) vs = voxels.size(1) visible = jt.zeros_like(voxels).int32() voxels, visible = voxelization_cuda.voxelize_sub3(faces, voxels, visible) sum_visible = visible.sum() while True: voxels, visible = voxelization_cuda.voxelize_sub4( faces, voxels, visible) if visible.sum() == sum_visible: break else: sum_visible = visible.sum() return 1 - visible
def get_pos_proposal_indexes(self, locations, box_regression, matched_idxes, targets): locations = jt.contrib.concat(locations, dim=0) pos_indexes_for_targets = [] for im in range(len(targets)): pos_indexes_for_targets_per_im = locations.new_ones( len(targets[im])).long() * -1 box_regression_im = [ box_regression[l][im].detach().view(4, -1).transpose(0, 1) * self.fpn_strides[l] for l in range(len(box_regression)) ] box_regression_im = jt.contrib.concat(box_regression_im, dim=0) for t_id in range(len(targets[im])): valid = matched_idxes[im] == t_id if valid.sum() == 0: continue valid_location = locations[valid] valid_regression = box_regression_im[valid] detections = jt.stack([ valid_location[:, 0] - valid_regression[:, 0], valid_location[:, 1] - valid_regression[:, 1], valid_location[:, 0] + valid_regression[:, 2], valid_location[:, 1] + valid_regression[:, 3], ], dim=1) detect_boxlist = BoxList(detections, targets[im].size, mode="xyxy") target_boxlist = BoxList(targets[im].bbox[t_id:t_id + 1], targets[im].size, mode="xyxy") match_quality_matrix = boxlist_iou(detect_boxlist, target_boxlist) pos_labels_per_target = jt.zeros_like(valid) iou_in_target = match_quality_matrix[:, 0] pos_in_target = (iou_in_target == iou_in_target.max()) pos_labels_per_target[valid] = pos_in_target pos_indexes_for_targets_per_im[ t_id] = pos_labels_per_target.nonzero()[0][0] pos_indexes_for_targets.append(pos_indexes_for_targets_per_im) return pos_indexes_for_targets
def remove_isolated_nodes(edge_index, edge_attr=None, num_nodes=None): r"""Removes the isolated nodes from the graph given by :attr:`edge_index` with optional edge attributes :attr:`edge_attr`. In addition, returns a mask of shape :obj:`[num_nodes]` to manually filter out isolated node features later on. Self-loops are preserved for non-isolated nodes. Args: edge_index (Var int32): The edge indices. edge_attr (Var, optional): Edge weights or multi-dimensional edge features. (default: :obj:`None`) num_nodes (int, optional): The number of nodes, *i.e.* :obj:`max_val + 1` of :attr:`edge_index`. (default: :obj:`None`) :rtype: (Var int32, Var, Var bool) """ num_nodes = maybe_num_nodes(edge_index, num_nodes) out = segregate_self_loops(edge_index, edge_attr) edge_index, edge_attr, loop_edge_index, loop_edge_attr = out mask = jt.zeros((num_nodes), dtype=Var.bool) mask[edge_index.view(-1)] = 1 assoc = jt.full((num_nodes, ), -1, dtype=Var.int32) assoc[mask] = jt.arange(mask.sum()) edge_index = assoc[edge_index] loop_mask = jt.zeros_like(mask) loop_mask[loop_edge_index[0]] = 1 loop_mask = loop_mask & mask loop_assoc = jt.full_like(assoc, -1) loop_assoc[loop_edge_index[0]] = jt.arange(loop_edge_index.size(1)) loop_idx = loop_assoc[loop_mask] loop_edge_index = assoc[loop_edge_index[:, loop_idx]] edge_index = jt.concat([edge_index, loop_edge_index], dim=1) if edge_attr is not None: loop_edge_attr = loop_edge_attr[loop_idx] edge_attr = jt.concat([edge_attr, loop_edge_attr], dim=0) return edge_index, edge_attr, mask
def decode(self, rel_codes, boxes): """ From a set of original boxes and encoded relative box offsets, get the decoded boxes. Arguments: rel_codes (Tensor): encoded boxes boxes (Tensor): reference boxes. """ boxes = boxes.cast(rel_codes.dtype) TO_REMOVE = 1 # TODO remove widths = boxes[:, 2] - boxes[:, 0] + TO_REMOVE heights = boxes[:, 3] - boxes[:, 1] + TO_REMOVE ctr_x = boxes[:, 0] + 0.5 * widths ctr_y = boxes[:, 1] + 0.5 * heights wx, wy, ww, wh = self.weights dx = rel_codes[:, 0::4] / wx dy = rel_codes[:, 1::4] / wy dw = rel_codes[:, 2::4] / ww dh = rel_codes[:, 3::4] / wh # Prevent sending too large values into torch.exp() dw = jt.clamp(dw, max_v=self.bbox_xform_clip) dh = jt.clamp(dh, max_v=self.bbox_xform_clip) pred_ctr_x = dx * widths.unsqueeze(-1) + ctr_x.unsqueeze(-1) pred_ctr_y = dy * heights.unsqueeze(-1) + ctr_y.unsqueeze(-1) pred_w = jt.exp(dw) * widths.unsqueeze(-1) pred_h = jt.exp(dh) * heights.unsqueeze(-1) pred_boxes = jt.zeros_like(rel_codes) # x1 pred_boxes[:, 0::4] = pred_ctr_x - 0.5 * pred_w # y1 pred_boxes[:, 1::4] = pred_ctr_y - 0.5 * pred_h # x2 (note: "- 1" is correct; don't be fooled by the asymmetry) pred_boxes[:, 2::4] = pred_ctr_x + 0.5 * pred_w - 1 # y2 (note: "- 1" is correct; don't be fooled by the asymmetry) pred_boxes[:, 3::4] = pred_ctr_y + 0.5 * pred_h - 1 return pred_boxes
def compute_loss(p, targets, model): # predictions, targets, model lcls, lbox, lobj = jt.zeros((1, )), jt.zeros((1, )), jt.zeros((1, )) tcls, tbox, indices, anchors = build_targets(p, targets, model) # targets h = model.hyp # hyperparameters # Define criteria BCEcls = nn.BCEWithLogitsLoss(pos_weight=jt.array( [h['cls_pw']])) # weight=model.class_weights) BCEobj = nn.BCEWithLogitsLoss(pos_weight=jt.array([h['obj_pw']])) # Class label smoothing https://arxiv.org/pdf/1902.04103.pdf eqn 3 cp, cn = smooth_BCE(eps=0.0) # Focal loss g = h['fl_gamma'] # focal loss gamma if g > 0: BCEcls, BCEobj = FocalLoss(BCEcls, g), FocalLoss(BCEobj, g) # Losses balance = [4.0, 1.0, 0.4, 0.1] # P3-P6 for i, pi in enumerate(p): # layer index, layer predictions b, a, gj, gi = indices[i] # image, anchor, gridy, gridx tobj = jt.zeros_like(pi[..., 0]) # target obj n = b.shape[0] # number of targets if n: ps = pi[b, a, gj, gi] # prediction subset corresponding to targets # Regression pxy = ps[:, :2].sigmoid() * 2. - 0.5 pwh = (ps[:, 2:4].sigmoid() * 2)**2 * anchors[i] pbox = jt.contrib.concat((pxy, pwh), 1) # predicted box iou = bbox_iou(pbox.transpose(1, 0), tbox[i], x1y1x2y2=False, CIoU=True) # iou(prediction, target) lbox += (1.0 - iou).mean() # iou loss # Objectness tobj[b, a, gj, gi] = (1.0 - model.gr) + model.gr * iou.detach().clamp(0).cast( tobj.dtype) # iou ratio # Classification if model.nc > 1: # cls loss (only if multiple classes) t = jt.full_like(ps[:, 5:], cn) # targets t[list(range(n)), tcls[i]] = cp lcls += BCEcls(ps[:, 5:], t) # BCE # Append targets to text file # with open('targets.txt', 'a') as file: # [file.write('%11.5g ' * 4 % tuple(x) + '\n') for x in jt.contrib.concat((txy[i], twh[i]), 1)] lobj += BCEobj(pi[..., 4], tobj) * balance[i] # obj loss lbox *= h['box'] lobj *= h['obj'] lcls *= h['cls'] bs = tobj.shape[0] # batch size loss = lbox + lobj + lcls return loss * bs, jt.contrib.concat((lbox, lobj, lcls, loss)).detach()
def build_targets(p, targets, model): # Build targets for compute_loss(), input targets(image,class,x,y,w,h) det = model.model[-1] # Detect() module na, nt = det.na, targets.shape[0] # number of anchors, targets tcls, tbox, indices, anch = [], [], [], [] gain = jt.ones((7, )) # normalized to gridspace gain ai = jt.index( (na, ), dim=0).float().view(na, 1).repeat(1, nt) # same as .repeat_interleave(nt) targets = jt.contrib.concat((targets.repeat(na, 1, 1), ai[:, :, None]), 2) # append anchor indices g = 0.5 # bias off = jt.array( [ [0, 0], # [1, 0], [0, 1], [-1, 0], [0, -1], # j,k,l,m # [1, 1], [1, -1], [-1, 1], [-1, -1], # jk,jm,lk,lm ], ).float() * g # offsets for i in range(det.nl): anchors = det.anchors[i] gain[2:6] = jt.array( [p[i].shape[3], p[i].shape[2], p[i].shape[3], p[i].shape[2]]) # xyxy gain # Match targets to anchors t = targets * gain if nt: # Matches r = t[:, :, 4:6] / anchors[:, None] # wh ratio j = jt.maximum(r, 1. / r).max(2) < model.hyp['anchor_t'] # compare # j = wh_iou(anchors, t[:, 4:6]) > model.hyp['iou_t'] # iou(3,n)=wh_iou(anchors(3,2), gwh(n,2)) t = t[j] # filter # Offsets gxy = t[:, 2:4] # grid xy gxi = gain[jt.array([2, 3])] - gxy # inverse # j, k = jt.logical_and((gxy % 1. < g), (gxy > 1.)).int().transpose(1,0).bool() # l, m = jt.logical_and((gxi % 1. < g),(gxi > 1.)).int().transpose(1,0).bool() jk = jt.logical_and((gxy % 1. < g), (gxy > 1.)) lm = jt.logical_and((gxi % 1. < g), (gxi > 1.)) j, k = jk[:, 0], jk[:, 1] l, m = lm[:, 0], lm[:, 1] j = jt.stack((jt.ones_like(j), )) t = t.repeat((off.shape[0], 1, 1))[j] offsets = (jt.zeros_like(gxy)[None] + off[:, None])[j] else: t = targets[0] offsets = 0 # Define b = t[:, 0].int32() c = t[:, 1].int32() # image, class gxy = t[:, 2:4] # grid xy gwh = t[:, 4:6] # grid wh gij = (gxy - offsets).int32() gi, gj = gij[:, 0], gij[:, 1] # grid xy indices # Append a = t[:, 6].int32() # anchor indices indices.append((b, a, gj.clamp(0, gain[3] - 1), gi.clamp(0, gain[2] - 1))) # image, anchor, grid indices tbox.append(jt.contrib.concat((gxy - gij, gwh), 1)) # box anch.append(anchors[a]) # anchors tcls.append(c) # class return tcls, tbox, indices, anch