def boxes_to_masks(boxes, h, w, padding=0.0): n = boxes.shape[0] boxes = boxes x1 = boxes[:, 0] y1 = boxes[:, 1] x2 = boxes[:, 2] y2 = boxes[:, 3] b_w = x2 - x1 b_h = y2 - y1 x1 = jt.clamp(x1 - 1 - b_w * padding, min_v=0) x2 = jt.clamp(x2 + 1 + b_w * padding, max_v=w) y1 = jt.clamp(y1 - 1 - b_h * padding, min_v=0) y2 = jt.clamp(y2 + 1 + b_h * padding, max_v=h) rows = jt.arange(w, dtype=x1.dtype).view(1, 1, -1).expand((n, h, w)) cols = jt.arange(h, dtype=x1.dtype).view(1, -1, 1).expand((n, h, w)) masks_left = rows >= x1.view(-1, 1, 1) masks_right = rows < x2.view(-1, 1, 1) masks_up = cols >= y1.view(-1, 1, 1) masks_down = cols < y2.view(-1, 1, 1) masks = masks_left * masks_right * masks_up * masks_down return masks
def crop(masks, boxes, padding: int = 1): """ "Crop" predicted masks by zeroing out everything not in the predicted bbox. Vectorized by Chong (thanks Chong). Args: - masks should be a size [h, w, n] tensor of masks - boxes should be a size [n, 4] tensor of bbox coords in relative point form """ h, w, n = masks.shape x1, x2 = sanitize_coordinates(boxes[:, 0], boxes[:, 2], w, padding, cast=False) y1, y2 = sanitize_coordinates(boxes[:, 1], boxes[:, 3], h, padding, cast=False) rows = jt.arange(w, dtype=x1.dtype).view(1, -1, 1).expand((h, w, n)) cols = jt.arange(h, dtype=x1.dtype).view(-1, 1, 1).expand((h, w, n)) masks_left = rows >= x1.view(1, 1, -1) masks_right = rows < x2.view(1, 1, -1) masks_up = cols >= y1.view(1, 1, -1) masks_down = cols < y2.view(1, 1, -1) crop_mask = masks_left * masks_right * masks_up * masks_down return masks * crop_mask.float()
def compute_locations_per_level(self, h, w, stride): shifts_x = jt.arange(0, w * stride, step=stride, dtype='float32') shifts_y = jt.arange(0, h * stride, step=stride, dtype='float32') shift_y, shift_x = jt.meshgrid(shifts_y, shifts_x) shift_x = shift_x.reshape(-1) shift_y = shift_y.reshape(-1) locations = jt.stack((shift_x, shift_y), dim=1) + stride // 2 return locations
def execute(self, x, boxes): """ Arguments: x (Tensor): the mask logits boxes (list[BoxList]): bounding boxes that are used as reference, one for each image Returns: results (list[BoxList]): one BoxList for each image, containing the extra field mask """ mask_prob = x.sigmoid() # select masks coresponding to the predicted classes num_masks = x.shape[0] labels = [bbox.get_field("labels") for bbox in boxes] labels = jt.contrib.concat(labels, dim=0) index = jt.arange(num_masks) mask_prob = mask_prob[index, labels].unsqueeze(1) boxes_per_image = [len(box) for box in boxes] mask_prob = mask_prob.split(boxes_per_image, dim=0) if self.masker: mask_prob = self.masker(mask_prob, boxes) results = [] for prob, box in zip(mask_prob, boxes): bbox = BoxList(box.bbox, box.size, mode="xyxy") for field in box.fields(): bbox.add_field(field, box.get_field(field)) bbox.add_field("mask", prob) results.append(bbox) return results
def __call__(self, proposals, mask_logits, targets): """ Arguments: proposals (list[BoxList]) mask_logits (Tensor) targets (list[BoxList]) Return: mask_loss (Tensor): scalar tensor containing the loss """ labels, mask_targets, mask_ratios = self.prepare_targets( proposals, targets) labels = cat(labels, dim=0) mask_targets = cat(mask_targets, dim=0) positive_inds = jt.nonzero(labels > 0).squeeze(1) labels_pos = labels[positive_inds] # accept empty tensors, so handle it separately if mask_targets.numel() == 0: if not self.maskiou_on: return mask_logits.sum() * 0 else: selected_index = jt.arange(mask_logits.shape[0]) selected_mask = mask_logits[selected_index, labels] mask_num, mask_h, mask_w = selected_mask.shape selected_mask = selected_mask.reshape(mask_num, 1, mask_h, mask_w) return mask_logits.sum() * 0, selected_mask, labels, None if self.maskiou_on: mask_ratios = cat(mask_ratios, dim=0) value_eps = 1e-10 * jt.ones((mask_targets.shape[0], )) mask_ratios = jt.maximum(mask_ratios, value_eps) pred_masks = mask_logits[positive_inds, labels_pos] pred_masks[:] = pred_masks > 0 mask_targets_full_area = mask_targets.sum( dims=[1, 2]) / mask_ratios mask_ovr = pred_masks * mask_targets mask_ovr_area = mask_ovr.sum(dims=[1, 2]) mask_union_area = pred_masks.sum( dims=[1, 2]) + mask_targets_full_area - mask_ovr_area value_1 = jt.ones((pred_masks.shape[0], )) value_0 = jt.zeros((pred_masks.shape[0], )) mask_union_area = jt.maximum(mask_union_area, value_1) mask_ovr_area = jt.maximum(mask_ovr_area, value_0) maskiou_targets = mask_ovr_area / mask_union_area binary_cross_entropy_with_logits = nn.BCEWithLogitsLoss() mask_loss = binary_cross_entropy_with_logits( mask_logits[positive_inds, labels_pos], mask_targets) if not self.maskiou_on: return mask_loss else: selected_index = jt.index((mask_logits.shape[0], ), dim=0) selected_mask = mask_logits[selected_index, labels] mask_num, mask_h, mask_w = selected_mask.shape selected_mask = selected_mask.reshape(mask_num, 1, mask_h, mask_w) selected_mask = selected_mask.sigmoid() return mask_loss, selected_mask, labels, maskiou_targets
def test_normal(self): from jittor import init n = 10000 r = 0.155 a = init.gauss([n], "float32", 1, 3) data = a.data assert (np.abs((data < (1 - 3)).mean() - r) < 0.1) assert (np.abs((data < (1)).mean() - 0.5) < 0.1) assert (np.abs((data < (1 + 3)).mean() - (1 - r)) < 0.1) np_res = np.random.normal(1, 0.1, (100, 100)) jt_res = jt.normal(1., 0.1, (100, 100)) assert (np.abs(np_res.mean() - jt_res.data.mean()) < 0.1) assert (np.abs(np_res.std() - jt_res.data.std()) < 0.1) np_res = torch.normal(torch.arange(1., 10000.), 1) jt_res = jt.normal(jt.arange(1, 10000), 1) assert (np.abs(np_res.mean() - jt_res.data.mean()) < 0.1) assert (np.abs(np_res.std() - jt_res.data.std()) < 1) np_res = np.random.randn(100, 100) jt_res = jt.randn(100, 100) assert (np.abs(np_res.mean() - jt_res.data.mean()) < 0.1) assert (np.abs(np_res.std() - jt_res.data.std()) < 0.1) np_res = np.random.rand(100, 100) jt_res = jt.rand(100, 100) assert (np.abs(np_res.mean() - jt_res.data.mean()) < 0.1) assert (np.abs(np_res.std() - jt_res.data.std()) < 0.1)
def add_self_loops(edge_index, edge_weight: Optional[Var] = None, fill_value: float = 1., num_nodes: Optional[int] = None): r"""Adds a self-loop :math:`(i,i) \in \mathcal{E}` to every node :math:`i \in \mathcal{V}` in the graph given by :attr:`edge_index`. In case the graph is weighted, self-loops will be added with edge weights denoted by :obj:`fill_value`. Args: edge_index (Var int32): The edge indices. edge_weight (Var, optional): One-dimensional edge weights. (default: :obj:`None`) fill_value (float, optional): If :obj:`edge_weight` is not :obj:`None`, will add self-loops with edge weights of :obj:`fill_value` to the graph. (default: :obj:`1.`) num_nodes (int, optional): The number of nodes, *i.e.* :obj:`max_val + 1` of :attr:`edge_index`. (default: :obj:`None`) :rtype: (:class:`Var int32`, :class:`Var`) """ N = maybe_num_nodes(edge_index, num_nodes) loop_index = jt.arange(0, N, dtype=Var.int32) loop_index = loop_index.unsqueeze(0).repeat(2, 1) if edge_weight is not None: assert edge_weight.numel() == edge_index.size(1) loop_weight = init.constant((N, ), edge_weight.dtype, fill_value) edge_weight = jt.concat([edge_weight, loop_weight], dim=0) edge_index = jt.concat([edge_index, loop_index], dim=1) return edge_index, edge_weight
def grid_anchors(self, grid_sizes): anchors = [] for size, stride, base_anchors in zip(grid_sizes, self.strides, self.cell_anchors): grid_height, grid_width = size shifts_x = jt.arange(0, grid_width * stride, step=stride).float32() shifts_y = jt.arange(0, grid_height * stride, step=stride).float32() shift_y, shift_x = jt.meshgrid(shifts_y, shifts_x) shift_x = shift_x.reshape(-1) shift_y = shift_y.reshape(-1) shifts = jt.stack((shift_x, shift_y, shift_x, shift_y), dim=1) anchors.append((shifts.reshape(-1, 1, 4) + base_anchors.reshape(1, -1, 4)).reshape(-1, 4)) return anchors
def execute(self, x): assert x.ndim == 3 # interp = torch.lerp(self.avg_latent, x, self.threshold) interp = self.avg_latent + self.threshold * (x - self.avg_latent) # do_trunc = (torch.arange(x.size(1)) < self.max_layer).view(1, -1, 1).to(x.device) do_trunc = (jt.arange(x.size(1)) < self.max_layer).view(1, -1, 1) # return torch.where(do_trunc, interp, x) return do_trunc * interp + (1 - do_trunc) * x
def fast_nms(self, boxes, masks, scores, iou_threshold: float = 0.5, top_k: int = 200, second_threshold: bool = False): idx, scores = scores.argsort(1, descending=True) idx = idx[:, :top_k] scores = scores[:, :top_k] num_classes, num_dets = idx.shape boxes = boxes[idx.view(-1)].view(num_classes, num_dets, 4) masks = masks[idx.view(-1)].view(num_classes, num_dets, -1) iou = jaccard(boxes, boxes) iou = iou.triu_(diagonal=1) iou_max = iou.max(dim=1) # Now just filter out the ones higher than the threshold keep = (iou_max <= iou_threshold) # We should also only keep detections over the confidence threshold, but at the cost of # maxing out your detection count for every image, you can just not do that. Because we # have such a minimal amount of computation per detection (matrix mulitplication only), # this increase doesn't affect us much (+0.2 mAP for 34 -> 33 fps), so we leave it out. # However, when you implement this in your method, you should do this second threshold. if second_threshold: keep *= (scores > self.conf_thresh) # Assign each kept detection to its corresponding class classes = jt.arange(num_classes).unsqueeze(1).expand_as(keep) classes = classes[keep] boxes = boxes[keep] masks = masks[keep] scores = scores[keep] #print('keep finish') # Only keep the top cfg.max_num_detections highest scores across all classes idx, scores = jt.argsort(scores, dim=0, descending=True) #print('argsort finish') idx = idx[:cfg.max_num_detections] scores = scores[:cfg.max_num_detections] classes = classes[idx] boxes = boxes[idx] masks = masks[idx] ''' scores = scores[:cfg.max_num_detections] classes = classes[:cfg.max_num_detections] boxes = boxes[:cfg.max_num_detections] masks = masks[:cfg.max_num_detections] ''' return boxes, masks, classes, scores
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 crop_by_box(masks, box, padding=0.0): n, h, w = masks.size() b_w = box[:, 2] - box[:, 0] b_h = box[:, 3] - box[:, 1] x1 = jt.clamp(box[:, 0] - b_w * padding - 1, min_v=0) x2 = jt.clamp(box[:, 2] + b_w * padding + 1, max_v=w - 1) y1 = jt.clamp(box[:, 1] - b_h * padding - 1, min_v=0) y2 = jt.clamp(box[:, 3] + b_h * padding + 1, max_v=h - 1) rows = jt.arange(w, dtype=x1.dtype).view(1, 1, -1).expand((n, h, w)) cols = jt.arange(h, dtype=x1.dtype).view(1, -1, 1).expand((n, h, w)) masks_left = rows >= x1.view(n, 1, 1) masks_right = rows < x2.view(n, 1, 1) masks_up = cols >= y1.view(n, 1, 1) masks_down = cols < y2.view(n, 1, 1) crop_mask = masks_left * masks_right * masks_up * masks_down return masks * crop_mask.float(), crop_mask
def sigmoid_focal_loss(logits, targets, gamma, alpha): num_classes = logits.shape[1] dtype = targets.dtype class_range = jt.arange(1, num_classes + 1, dtype=dtype).unsqueeze(0) t = targets.unsqueeze(1) p = logits.sigmoid() term1 = (1 - p)**gamma * jt.log(p) term2 = p**gamma * jt.log(1 - p) return -(t == class_range).float() * term1 * alpha - ( (t != class_range) * (t >= 0)).float() * term2 * (1 - alpha)
def traditional_nms(self, boxes, masks, scores, iou_threshold=0.5, conf_thresh=0.05): import pyximport pyximport.install(setup_args={"include_dirs": np.get_include()}, reload_support=True) from utils.cython_nms import nms as cnms num_classes = scores.shape[0] idx_lst = [] cls_lst = [] scr_lst = [] # Multiplying by max_size is necessary because of how cnms computes its area and intersections boxes = boxes * cfg.max_size for _cls in range(num_classes): cls_scores = scores[_cls] conf_mask = cls_scores > conf_thresh idx = jt.arange(cls_scores.shape[0], device=boxes.device) cls_scores = cls_scores[conf_mask] idx = idx[conf_mask] if cls_scores.shape[0] == 0: continue preds = jt.contrib.concat([boxes[conf_mask], cls_scores[:, None]], dim=1).numpy() keep = cnms(preds, iou_threshold) keep = jt.array(keep).int64() idx_lst.append(idx[keep]) cls_lst.append(keep * 0 + _cls) scr_lst.append(cls_scores[keep]) idx = jt.contrib.concat(idx_lst, dim=0) classes = jt.contrib.concat(cls_lst, dim=0) scores = jt.contrib.concat(scr_lst, dim=0) scores, idx2 = scores.sort(0, descending=True) idx2 = idx2[:cfg.max_num_detections] scores = scores[:cfg.max_num_detections] idx = idx[idx2] classes = classes[idx2] # Undo the multiplication above return boxes[idx] / cfg.max_size, masks[idx], classes, scores
def random_subsample(pcd, n_points=2048): """ Args: pcd: (B, N, 3) returns: new_pcd: (B, n_points, 3) """ b, n, _ = pcd.shape batch_idx = jittor.arange(b).reshape((-1, 1)).repeat(1, n_points) idx = jittor.concat([jittor.randperm(n)[:n_points].reshape((1, -1)) for i in range(b)], 0) return pcd[batch_idx, idx, :]
def grad(self, grad): x, numangle, numrho = self.save_vars cuda_src_backward = csb.replace('#numangle', str(numangle)) cuda_src_backward = cuda_src_backward.replace('#numrho', str(numrho)) irho = int((h * h + w * w)**0.5 + 1) / float((numrho - 1)) itheta = 3.14159265358979323846 / numangle angle = jt.arange(numangle) * itheta tabCos = angle.cos() / irho tabSin = angle.sin() / irho return jt.code([x.shape], [x.dtype], [x, grad, tabCos, tabSin], cuda_src=cuda_src_backward)
def _forward_train(self,features,img_size,boxes,labels): N = features.shape[0] rpn_locs, rpn_scores, rois, roi_indices, anchor = self.rpn(features, img_size) sample_rois = [] gt_roi_locs = [] gt_roi_labels = [] sample_roi_indexs = [] gt_rpn_locs = [] gt_rpn_labels = [] for i in range(N): index = jt.where(roi_indices == i)[0] roi = rois[index,:] box = boxes[i] label = labels[i] sample_roi, gt_roi_loc, gt_roi_label = self.proposal_target_creator(roi,box,label) sample_roi_index = i*jt.ones((sample_roi.shape[0],)) sample_rois.append(sample_roi) gt_roi_labels.append(gt_roi_label) gt_roi_locs.append(gt_roi_loc) sample_roi_indexs.append(sample_roi_index) gt_rpn_loc, gt_rpn_label = self.anchor_target_creator(box,anchor,img_size) gt_rpn_locs.append(gt_rpn_loc) gt_rpn_labels.append(gt_rpn_label) sample_roi_indexs = jt.contrib.concat(sample_roi_indexs,dim=0) sample_rois = jt.contrib.concat(sample_rois,dim=0) roi_cls_loc, roi_score = self.head(features,sample_rois,sample_roi_indexs) # ------------------ RPN losses -------------------# rpn_locs = rpn_locs.reshape(-1,4) rpn_scores = rpn_scores.reshape(-1,2) gt_rpn_labels = jt.contrib.concat(gt_rpn_labels,dim=0) gt_rpn_locs = jt.contrib.concat(gt_rpn_locs,dim=0) rpn_loc_loss = _fast_rcnn_loc_loss(rpn_locs,gt_rpn_locs,gt_rpn_labels,self.rpn_sigma) rpn_cls_loss = nn.cross_entropy_loss(rpn_scores[gt_rpn_labels>=0,:],gt_rpn_labels[gt_rpn_labels>=0]) # ------------------ ROI losses (fast rcnn loss) -------------------# gt_roi_locs = jt.contrib.concat(gt_roi_locs,dim=0) gt_roi_labels = jt.contrib.concat(gt_roi_labels,dim=0) n_sample = roi_cls_loc.shape[0] roi_cls_loc = roi_cls_loc.view(n_sample, np.prod(roi_cls_loc.shape[1:]).item()//4, 4) roi_loc = roi_cls_loc[jt.arange(0, n_sample).int32(), gt_roi_labels] roi_loc_loss = _fast_rcnn_loc_loss(roi_loc,gt_roi_locs,gt_roi_labels,self.roi_sigma) roi_cls_loss = nn.cross_entropy_loss(roi_score, gt_roi_labels) losses = [rpn_loc_loss, rpn_cls_loss, roi_loc_loss, roi_cls_loss] losses = losses + [sum(losses)] return losses
def forward_for_single_feature_map(self, anchors, objectness, box_regression): """ Arguments: anchors: list[BoxList] objectness: tensor of size N, A, H, W box_regression: tensor of size N, A * 4, H, W """ N, A, H, W = objectness.shape # put in the same format as anchors objectness = permute_and_flatten(objectness, N, A, 1, H, W).reshape(N, -1) objectness = objectness.sigmoid() box_regression = permute_and_flatten(box_regression, N, A, 4, H, W) num_anchors = A * H * W pre_nms_top_n = min(self.pre_nms_top_n, num_anchors) objectness, topk_idx = objectness.topk(pre_nms_top_n, dim=1, sorted=True) batch_idx = jt.arange(N).unsqueeze(1) box_regression = box_regression[batch_idx,topk_idx] image_shapes = [box.size for box in anchors] concat_anchors = jt.contrib.concat([a.bbox for a in anchors], dim=0) concat_anchors = concat_anchors.reshape(N, -1, 4)[batch_idx, topk_idx] proposals = self.box_coder.decode( box_regression.reshape(-1, 4), concat_anchors.reshape(-1, 4) ) proposals = proposals.reshape(N, -1, 4) result = [] for i in range(len(image_shapes)): proposal, score, im_shape = proposals[i], objectness[i], image_shapes[i] boxlist = BoxList(proposal, im_shape, mode="xyxy") boxlist.add_field("objectness", score) boxlist = boxlist.clip_to_image(remove_empty=False) boxlist = remove_small_boxes(boxlist, self.min_size) boxlist = boxlist_nms( boxlist, self.nms_thresh, max_proposals=self.post_nms_top_n, score_field="objectness", ) result.append(boxlist) return result
def index2d(src, idx): """ Indexes a tensor by a 2d index. In effect, this does out[i, j] = src[i, idx[i, j]] Both src and idx should have the same size. """ offs = jt.arange((idx.shape[0])).unsqueeze(1).expand_as(idx) idx = idx + offs * idx.shape[1] return src.view(-1)[idx.view(-1)].view(idx.shape)
def execute(self, x, numangle, numrho): n, c, h, w = x.shape cuda_src_forward = csf.replace('#numangle', str(numangle)) cuda_src_forward = cuda_src_forward.replace('#numrho', str(numrho)) irho = int((h * h + w * w)**0.5 + 1) / float((numrho - 1)) itheta = 3.14159265358979323846 / numangle angle = jt.arange(numangle) * itheta tabCos = angle.cos() / irho tabSin = angle.sin() / irho output = jt.code([n, c, numangle, numrho], x.dtype, [x, tabCos, tabSin], cuda_src=cuda_src_forward) self.save_vars = x, numangle, numrho return output
def add_remaining_self_loops(edge_index, edge_weight: Optional[Var] = None, fill_value: float = 1., num_nodes: Optional[int] = None): r"""Adds remaining self-loop :math:`(i,i) \in \mathcal{E}` to every node :math:`i \in \mathcal{V}` in the graph given by :attr:`edge_index`. In case the graph is weighted and already contains a few self-loops, only non-existent self-loops will be added with edge weights denoted by :obj:`fill_value`. Args: edge_index (Var int32): The edge indices. edge_weight (Var, optional): One-dimensional edge weights. (default: :obj:`None`) fill_value (float, optional): If :obj:`edge_weight` is not :obj:`None`, will add self-loops with edge weights of :obj:`fill_value` to the graph. (default: :obj:`1.`) num_nodes (int, optional): The number of nodes, *i.e.* :obj:`max_val + 1` of :attr:`edge_index`. (default: :obj:`None`) :rtype: (:class:`Var int32`, :class:`Var`) """ N = maybe_num_nodes(edge_index, num_nodes) row, col = edge_index[0], edge_index[1] mask = row != col loop_index = jt.arange(0, N, dtype=row.dtype) loop_index = loop_index.unsqueeze(0).repeat(2, 1) edge_index = jt.concat([edge_index[:, mask], loop_index], dim=1) if edge_weight is not None: inv_mask = jt.logical_not(mask) loop_weight = init.constant((N, ), edge_weight.dtype, fill_value) remaining_edge_weight = edge_weight[inv_mask] if remaining_edge_weight.numel() > 0: loop_weight[row[inv_mask]] = remaining_edge_weight edge_weight = jt.concat([edge_weight[mask], loop_weight], dim=0) return edge_index, edge_weight
def test_scatter(self): src = jt.arange(1, 11).reshape((2, 5)) index = jt.array([[0, 1, 2, 0]]) x = jt.zeros((3, 5), dtype=src.dtype).scatter_(0, index, src) assert (x.data == [[1, 0, 0, 4, 0], [0, 2, 0, 0, 0], [0, 0, 3, 0, 0]]).all() index = jt.array([[0, 1, 2], [0, 1, 4]]) x = jt.zeros((3, 5), dtype=src.dtype).scatter_(1, index, src) assert (x.data == [[1, 2, 3, 0, 0], [6, 7, 0, 0, 8], [0, 0, 0, 0, 0]]).all() x = jt.full((2, 4), 2.).scatter_(1, jt.array([[2], [3]]), jt.array(1.23), reduce='multiply') assert np.allclose(x.data, [[2.0000, 2.0000, 2.4600, 2.0000], [2.0000, 2.0000, 2.0000, 2.4600]]), x x = jt.full((2, 4), 2.).scatter_(1, jt.array([[2], [3]]), jt.array(1.23), reduce='add') assert np.allclose(x.data, [[2.0000, 2.0000, 3.2300, 2.0000], [2.0000, 2.0000, 2.0000, 3.2300]])
def vertex_normals(self): if self._vertex_normals_update: bs, nv = self.vertices.shape[:2] bs, nf = self.faces.shape[:2] faces = self.faces + (jt.arange(bs) * nv)[:, None, None] vertices_faces = self.vertices.reshape((bs * nv, 3))[faces.long()] faces = faces.view(-1, 3) vertices_faces = vertices_faces.view(-1, 3, 3) normals = (jt.cross( vertices_faces[:, 2] - vertices_faces[:, 1], vertices_faces[:, 0] - vertices_faces[:, 1])).reindex_reduce( op="sum", shape=[bs * nv, 3], indexes=["@e0(i0)", "i1"], extras=[faces[:, 1].long()]) + (jt.cross( vertices_faces[:, 0] - vertices_faces[:, 2], vertices_faces[:, 1] - vertices_faces[:, 2])).reindex_reduce( op="sum", shape=[bs * nv, 3], indexes=["@e0(i0)", "i1"], extras=[faces[:, 2].long()]) + (jt.cross( vertices_faces[:, 1] - vertices_faces[:, 0], vertices_faces[:, 2] - vertices_faces[:, 0])).reindex_reduce( op="sum", shape=[bs * nv, 3], indexes=["@e0(i0)", "i1"], extras=[faces[:, 0].long()]) normals = jt.normalize(normals, p=2, eps=1e-6, dim=1) self._vertex_normals = normals.reshape((bs, nv, 3)) self._vertex_normals_update = False return self._vertex_normals
def sample_images(batches_done): """Saves a generated sample from the validation set""" batches = next(iter(val_dataloader)) real_A = batches[0] real_A_eyel = batches[1] real_A_eyer = batches[2] real_A_nose = batches[3] real_A_mouth = batches[4] real_A_hair = batches[5] real_A_bg = batches[6] mask = batches[7] mask2 = batches[8] center = batches[9] cmaskel = batches[10] cmasker = batches[11] cmaskno = batches[12] cmaskmo = batches[13] maskh = mask * mask2 maskb = inverse_mask(mask2) fake_B0 = G_global(real_A) # EYES, NOSE, MOUTH fake_B_eyel1 = G_l_eyel(real_A_eyel) fake_B_eyel2 = ae_eyel(fake_B_eyel1) fake_B_eyel = add_with_mask(fake_B_eyel2, fake_B_eyel1, cmaskel) fake_B_eyer1 = G_l_eyer(real_A_eyer) fake_B_eyer2 = ae_eyer(fake_B_eyer1) fake_B_eyer = add_with_mask(fake_B_eyer2, fake_B_eyer1, cmasker) fake_B_nose1 = G_l_nose(real_A_nose) fake_B_nose2 = ae_nose(fake_B_nose1) fake_B_nose = add_with_mask(fake_B_nose2, fake_B_nose1, cmaskno) fake_B_mouth1 = G_l_mouth(real_A_mouth) outputs1 = CLm(real_A_mouth) pred = jt.argmax(outputs1, dim=1)[0] fake_B_mouth2w = ae_mowhite(fake_B_mouth1) fake_B_mouth2b = ae_moblack(fake_B_mouth1) fake_B_mouth2s = jt.contrib.concat((fake_B_mouth2w, fake_B_mouth2b), 1) bs, c, h, w = fake_B_mouth2s.shape index = pred + jt.arange(bs) * c fake_B_mouth2 = fake_B_mouth2s.reshape([-1, h, w])[index].reshape([bs, 1, h, w]) fake_B_mouth = add_with_mask(fake_B_mouth2, fake_B_mouth1, cmaskmo) # HAIR & BG outputs2 = CLh(real_A_hair) onehot2 = getonehot(outputs2, 3, bs) fake_B_hair = G_l_hair(real_A_hair, onehot2) fake_B_bg = G_l_bg(real_A_bg) # PARTCOMBINE fake_B1 = partCombiner2_bg(center, fake_B_eyel, fake_B_eyer, fake_B_nose, fake_B_mouth, fake_B_hair, fake_B_bg, maskh, maskb, comb_op=1, load_h=opt.img_height, load_w=opt.img_width) # FUSION NET fake_B = G_combine(jt.contrib.concat((fake_B0, fake_B1), 1)) img_sample = np.concatenate( [real_A.data, fake_B.repeat(1, 3, 1, 1).data], -2) save_image(img_sample, "images/%s/%s.jpg" % (opt.dataset_name, batches_done), nrow=5)
fake_B_eyel2 = ae_eyel(fake_B_eyel1) fake_B_eyel = add_with_mask(fake_B_eyel2, fake_B_eyel1, cmaskel) fake_B_eyer1 = G_l_eyer(real_A_eyer) fake_B_eyer2 = ae_eyer(fake_B_eyer1) fake_B_eyer = add_with_mask(fake_B_eyer2, fake_B_eyer1, cmasker) fake_B_nose1 = G_l_nose(real_A_nose) fake_B_nose2 = ae_nose(fake_B_nose1) fake_B_nose = add_with_mask(fake_B_nose2, fake_B_nose1, cmaskno) fake_B_mouth1 = G_l_mouth(real_A_mouth) outputs1 = CLm(real_A_mouth) pred = jt.argmax(outputs1, dim=1)[0] fake_B_mouth2w = ae_mowhite(fake_B_mouth1) fake_B_mouth2b = ae_moblack(fake_B_mouth1) fake_B_mouth2s = jt.contrib.concat((fake_B_mouth2w, fake_B_mouth2b), 1) bs, c, h, w = fake_B_mouth2s.shape index = pred + jt.arange(bs) * c fake_B_mouth2 = fake_B_mouth2s.reshape([-1, h, w])[index].reshape( [bs, 1, h, w]) fake_B_mouth = add_with_mask(fake_B_mouth2, fake_B_mouth1, cmaskmo) # HAIR & BG outputs2 = CLh(real_A_hair) onehot2 = getonehot(outputs2, 3, bs) fake_B_hair = G_l_hair(real_A_hair, onehot2) fake_B_bg = G_l_bg(real_A_bg) # PARTCOMBINE fake_B1 = partCombiner2_bg(center, fake_B_eyel, fake_B_eyer, fake_B_nose, fake_B_mouth, fake_B_hair,
def read_planetoid_data(folder, prefix): names = ['x', 'tx', 'allx', 'y', 'ty', 'ally', 'graph', 'test.index'] items = [read_file(folder, prefix, name) for name in names] x, tx, allx, y, ty, ally, graph, test_index = items train_index = jt.arange(y.size(0), dtype=Var.int32) val_index = jt.arange(y.size(0), y.size(0) + 500, dtype=Var.int32) sorted_test_index = test_index.argsort()[1] if prefix.lower() == 'citeseer': # There are some isolated nodes in the Citeseer graph, resulting in # none consecutive test indices. We need to identify them and add them # as zero vectors to `tx` and `ty`. len_test_indices = (test_index.max() - test_index.min()).item() + 1 tx_ext = jt.zeros((len_test_indices, tx.size(1))) tx_ext[sorted_test_index - test_index.min(), :] = tx ty_ext = jt.zeros((len_test_indices, ty.size(1))) ty_ext[sorted_test_index - test_index.min(), :] = ty tx, ty = tx_ext, ty_ext if prefix.lower() == 'nell.0.001': tx_ext = jt.zeros((len(graph) - allx.size(0), x.size(1))) tx_ext[sorted_test_index - allx.size(0)] = tx ty_ext = jt.zeros((len(graph) - ally.size(0), y.size(1))) ty_ext[sorted_test_index - ally.size(0)] = ty tx, ty = tx_ext, ty_ext x = jt.concat([allx, tx], dim=0) x[test_index] = x[sorted_test_index] # Creating feature vectors for relations. row, col, value = SparseTensor.from_dense(x).coo() rows, cols, values = [row], [col], [value] mask1 = index_to_mask(test_index, size=len(graph)) mask2 = index_to_mask(jt.arange(allx.size(0), len(graph)), size=len(graph)) mask = jt.logical_or(jt.logical_not(mask1), jt.logical_not(mask2)) isolated_index = mask.nonzero(as_tuple=False).view(-1)[allx.size(0):] rows += [isolated_index] cols += [jt.arange(isolated_index.size(0)) + x.size(1)] values += [jt.ones((isolated_index.size(0)))] x = SparseTensor(row=jt.concat(rows), col=jt.concat(cols), value=jt.concat(values)) else: x = jt.concat([allx, tx], dim=0) x[test_index] = x[sorted_test_index] y = jt.concat([ally, ty], dim=0).argmax(dim=1)[0] y[test_index] = y[sorted_test_index] train_mask = index_to_mask(train_index, size=y.size(0)) val_mask = index_to_mask(val_index, size=y.size(0)) test_mask = index_to_mask(test_index, size=y.size(0)) edge_index = edge_index_from_dict(graph, num_nodes=y.size(0)) data = Data(x=x, edge_index=edge_index, y=y) data.train_mask = train_mask data.val_mask = val_mask data.test_mask = test_mask return data
def __init__(self, n, weight=False): # n: number of inputs super(Sum, self).__init__() self.weight = weight # apply weights boolean self.iter = range(n - 1) # iter object if weight: self.w = -jt.arange(1., n) / 2 # layer weights
def test_ellipsis_with_none(self): a = jt.arange(2 * 4 * 4).reshape(2, 4, 4) b = a[..., :, None, :2] assert b.shape == [2, 4, 1, 2] np.testing.assert_allclose(b.data, a.data[..., :, None, :2])
def forward_for_single_feature_map(self, anchors, objectness, box_regression): """ Arguments: anchors: list[BoxList] objectness: tensor of size N, A, H, W box_regression: tensor of size N, A * 4, H, W """ # global II # import pickle N, A, H, W = objectness.shape # put in the same format as anchors objectness = permute_and_flatten(objectness, N, A, 1, H, W).reshape(N, -1) # print('objectness',objectness.mean()) objectness = objectness.sigmoid() box_regression = permute_and_flatten(box_regression, N, A, 4, H, W) # print('regression',box_regression.mean()) num_anchors = A * H * W pre_nms_top_n = min(self.pre_nms_top_n, num_anchors) # print(pre_nms_top_n) #print('objectness',objectness) # objectness = jt.array(pickle.load(open(f'/home/lxl/objectness_0_{II}.pkl','rb'))) # print(objectness.shape) objectness, topk_idx = objectness.topk(pre_nms_top_n, dim=1, sorted=True) # print(II,'topk',topk_idx.sum(),topk_idx.shape) batch_idx = jt.arange(N).unsqueeze(1) # pickle.dump(topk_idx.numpy(),open(f'/home/lxl/topk_idx_{II}_jt.pkl','wb')) # topk_idx_tmp = topk_idx.numpy() # batch_idx = jt.array(pickle.load(open(f'/home/lxl/batch_idx_{II}.pkl','rb'))) # topk_idx = jt.array(pickle.load(open(f'/home/lxl/topk_idx_{II}.pkl','rb'))) # err = np.abs(topk_idx_tmp-topk_idx.numpy()) # print('Error!!!!!!!!!!!!!!!!',err.sum()) # print(err.nonzero()) #print('box_regression0',box_regression) #batch_idx = jt.index(topk_idx.shape,dim=0) box_regression = box_regression[batch_idx, topk_idx] #print('box_regression1',box_regression) image_shapes = [box.size for box in anchors] concat_anchors = jt.contrib.concat([a.bbox for a in anchors], dim=0) concat_anchors = concat_anchors.reshape(N, -1, 4)[batch_idx, topk_idx] # box_regression = jt.array(pickle.load(open(f'/home/lxl/box_regression_{II}.pkl','rb'))) # concat_anchors = jt.array(pickle.load(open(f'/home/lxl/concat_anchors_{II}.pkl','rb'))) proposals = self.box_coder.decode(box_regression.reshape(-1, 4), concat_anchors.reshape(-1, 4)) proposals = proposals.reshape(N, -1, 4) # proposals = jt.array(pickle.load(open(f'/home/lxl/proposal_{II}.pkl','rb'))) # objectness = jt.array(pickle.load(open(f'/home/lxl/objectness_{II}.pkl','rb'))) # II+=1 result = [] for i in range(len(image_shapes)): proposal, score, im_shape = proposals[i], objectness[ i], image_shapes[i] boxlist = BoxList(proposal, im_shape, mode="xyxy") boxlist.add_field("objectness", score) boxlist = boxlist.clip_to_image(remove_empty=False) boxlist = remove_small_boxes(boxlist, self.min_size) boxlist = boxlist_nms( boxlist, self.nms_thresh, max_proposals=self.post_nms_top_n, score_field="objectness", ) result.append(boxlist) return result