def test_topk(descending, sorted, inp1d, kth_only): k = 3 if inp1d: data = np.random.permutation(7) else: data = np.random.permutation(5 * 7).reshape(5, 7) data = data.astype(np.int32) def np_sort(x): if descending: return np.sort(x)[..., ::-1] return np.sort(x) res = F.topk( tensor(data), k, descending=descending, no_sort=(not sorted), kth_only=kth_only ) values, indices = res values = values.numpy() indices = indices.numpy() if kth_only: np.testing.assert_equal( values, np.take_along_axis(data, indices[..., None], -1).squeeze(-1) ) np.testing.assert_equal(values, np_sort(data)[..., k - 1]) else: np.testing.assert_equal(values, np.take_along_axis(data, indices, -1)) if not sorted: values = np_sort(values) np.testing.assert_equal(values, np_sort(data)[..., :k])
def f(a, b): base = 0 c = b - a _, idx = F.topk(c, 3) # internally, biased_idx will be idx as gopt will ignore the addition biased_idx = base + idx return biased_idx
def find_top_rpn_proposals(self, rpn_cls_score_list, rpn_bbox_offset_list, anchors_list, im_info): prev_nms_top_n = (self.cfg.train_prev_nms_top_n if self.training else self.cfg.test_prev_nms_top_n) post_nms_top_n = (self.cfg.train_post_nms_top_n if self.training else self.cfg.test_post_nms_top_n) return_rois = [] for bid in range(im_info.shape[0]): batch_proposal_list = [] batch_score_list = [] batch_level_list = [] for l, (rpn_cls_score, rpn_bbox_offset, anchors) in enumerate( zip(rpn_cls_score_list, rpn_bbox_offset_list, anchors_list)): # get proposals and scores offsets = rpn_bbox_offset[bid].transpose(2, 3, 0, 1).reshape(-1, 4) proposals = self.box_coder.decode(anchors, offsets) scores = rpn_cls_score[bid].transpose(1, 2, 0).flatten() scores.detach() # prev nms top n scores, order = F.topk(scores, descending=True, k=prev_nms_top_n) proposals = proposals[order, :] batch_proposal_list.append(proposals) batch_score_list.append(scores) batch_level_list.append(F.full_like(scores, l)) # gather proposals, scores, level proposals = F.concat(batch_proposal_list, axis=0) scores = F.concat(batch_score_list, axis=0) levels = F.concat(batch_level_list, axis=0) proposals = layers.get_clipped_boxes(proposals, im_info[bid]) # filter invalid proposals and apply total level nms keep_mask = layers.filter_boxes(proposals) _, keep_inds = F.cond_take(keep_mask == 1, keep_mask) proposals = proposals[keep_inds, :] scores = scores[keep_inds] levels = levels[keep_inds] nms_keep_inds = layers.batched_nms(proposals, scores, levels, self.cfg.rpn_nms_threshold, post_nms_top_n) # generate rois to rcnn head, rois shape (N, 5), info [batch_id, x1, y1, x2, y2] rois = F.concat([proposals, scores.reshape(-1, 1)], axis=1) rois = rois[nms_keep_inds] batch_inds = F.full((rois.shape[0], 1), bid) batch_rois = F.concat([batch_inds, rois[:, :4]], axis=1) return_rois.append(batch_rois) return_rois = F.concat(return_rois, axis=0) return return_rois.detach()
def main(): parser = argparse.ArgumentParser() parser.add_argument("-a", "--arch", default="shufflenet_v2_x1_0", type=str) parser.add_argument("-m", "--model", default=None, type=str) parser.add_argument("-i", "--image", default=None, type=str) args = parser.parse_args() model = snet_model.__dict__[args.arch](pretrained=(args.model is None)) if args.model is not None: logging.info("load from checkpoint %s", args.model) checkpoint = megengine.load(args.model) if "state_dict" in checkpoint: state_dict = checkpoint["state_dict"] model.load_state_dict(state_dict) if args.image is None: path = "../../../assets/cat.jpg" else: path = args.image image = cv2.imread(path, cv2.IMREAD_COLOR) transform = T.Compose([ T.Resize(256), T.CenterCrop(224), T.Normalize(mean=[103.530, 116.280, 123.675], std=[57.375, 57.120, 58.395]), # BGR T.ToMode("CHW"), ]) def infer_func(processed_img): model.eval() logits = model(processed_img) probs = F.softmax(logits) return probs processed_img = transform.apply(image)[np.newaxis, :] probs = infer_func(processed_img) top_probs, classes = F.topk(probs, k=5, descending=True) with open("../../../assets/imagenet_class_info.json") as fp: imagenet_class_index = json.load(fp) for rank, (prob, classid) in enumerate( zip(top_probs.numpy().reshape(-1), classes.numpy().reshape(-1))): print("{}: class = {:20s} with probability = {:4.1f} %".format( rank, imagenet_class_index[str(classid)][1], 100 * prob))
def sample_labels(labels, num_samples, label_value, ignore_label=-1): """sample N labels with label value = sample_labels Args: labels(Tensor): shape of label is (N,) num_samples(int): label_value(int): Returns: label(Tensor): label after sampling """ assert labels.ndim == 1, "Only tensor of dim 1 is supported." mask = (labels == label_value) num_valid = mask.sum() if num_valid <= num_samples: return labels random_tensor = F.zeros_like(labels).astype("float32") random_tensor[mask] = uniform(size=num_valid) _, invalid_inds = F.topk(random_tensor, k=num_samples - num_valid) labels[invalid_inds] = ignore_label return labels
def sample_mask_from_labels(labels, num_sample, sample_value): """generate mask for labels using sampling method. Args: labels (Tensor): num_sample (int): sample_value (int): Returns: sample_mask (Tensor) """ assert labels.ndim == 1, "Only tensor of dim 1 is supported." # TODO: support bool mask sample_mask = (labels == sample_value).astype("float32") num_mask = sample_mask.sum().astype("int32") if num_mask <= num_sample: return sample_mask random_tensor = sample_mask * uniform(size=labels.shape) _, sampled_idx = F.topk(random_tensor, k=num_sample - num_mask) sample_mask[sampled_idx] = F.zeros(sampled_idx.shape) return sample_mask
def f(x): y = F.topk(x, 3) np.testing.assert_equal(y[0].shape.numpy(), np.array([ 3, ])) return y
def fwd(x): top, indices = F.topk(x, 5) return top, indices
def get_ground_truth(self, anchors_list, batched_gt_boxes, batched_num_gts): labels_list = [] offsets_list = [] ctrness_list = [] all_level_anchors = F.concat(anchors_list, axis=0) for bid in range(batched_gt_boxes.shape[0]): gt_boxes = batched_gt_boxes[bid, :batched_num_gts[bid]] ious = [] candidate_idxs = [] base = 0 for stride, anchors_i in zip(self.cfg.stride, anchors_list): ious.append( layers.get_iou( gt_boxes[:, :4], F.concat([ anchors_i - stride * self.cfg.anchor_scale / 2, anchors_i + stride * self.cfg.anchor_scale / 2, ], axis=1))) gt_centers = (gt_boxes[:, :2] + gt_boxes[:, 2:4]) / 2 distances = F.sqrt( F.sum((F.expand_dims(gt_centers, axis=1) - anchors_i)**2, axis=2)) _, topk_idxs = F.topk(distances, self.cfg.anchor_topk) candidate_idxs.append(base + topk_idxs) base += anchors_i.shape[0] ious = F.concat(ious, axis=1) candidate_idxs = F.concat(candidate_idxs, axis=1) candidate_ious = F.gather(ious, 1, candidate_idxs) ious_thr = (F.mean(candidate_ious, axis=1, keepdims=True) + F.std(candidate_ious, axis=1, keepdims=True)) is_foreground = F.scatter( F.zeros(ious.shape), 1, candidate_idxs, F.ones(candidate_idxs.shape)).astype(bool) & (ious >= ious_thr) is_in_boxes = F.min(self.point_coder.encode( all_level_anchors, F.expand_dims(gt_boxes[:, :4], axis=1)), axis=2) > 0 ious[~is_foreground] = -1 ious[~is_in_boxes] = -1 match_indices = F.argmax(ious, axis=0) gt_boxes_matched = gt_boxes[match_indices] anchor_max_iou = F.indexing_one_hot(ious, match_indices, axis=0) labels = gt_boxes_matched[:, 4].astype(np.int32) labels[anchor_max_iou == -1] = 0 offsets = self.point_coder.encode(all_level_anchors, gt_boxes_matched[:, :4]) left_right = offsets[:, [0, 2]] top_bottom = offsets[:, [1, 3]] ctrness = F.sqrt( F.clip(F.min(left_right, axis=1) / F.max(left_right, axis=1), lower=0) * F.clip(F.min(top_bottom, axis=1) / F.max(top_bottom, axis=1), lower=0)) labels_list.append(labels) offsets_list.append(offsets) ctrness_list.append(ctrness) return ( F.stack(labels_list, axis=0).detach(), F.stack(offsets_list, axis=0).detach(), F.stack(ctrness_list, axis=0).detach(), )
[(64, 512, 16, 16), (64, 512, 16, 16)], False, 10000, ), ( "subtensor", lambda x: x[0:20, 10:60], lambda x: x[0:20, 10:60], [(100, 100)], [(64, 512, 16, 16)], True, 1000, ), ( "topk", lambda x: MF.topk(x, 10), lambda x: torch.topk(x, 10), [(100, 100)], [(1000, 1000)], True, 1000, ), ( "tile", lambda x: MF.tile(x, (2, ) * len(x.shape)), lambda x: torch.tile(x, (2, ) * len(x.shape)), [(100, 100)], [(64, 512, 16, 16)], True, 1000, ),
def get_losses(self, anchors, pred_logits, pred_offsets, gt_boxes, im_info): # pylint: disable=too-many-statements def positive_bag_loss(logits, axis=1): weight = 1.0 / (1.0 - logits) weight /= weight.sum(axis=axis, keepdims=True) bag_prob = (weight * logits).sum(axis=1) return -layers.safelog(bag_prob) def negative_bag_loss(logits, gamma): return (logits**gamma) * (-layers.safelog(1.0 - logits)) pred_scores = F.sigmoid(pred_logits) box_prob_list = [] positive_losses = [] clamp_eps = 1e-7 bucket_size = self.cfg.bucket_size for bid in range(im_info.shape[0]): boxes_info = gt_boxes[bid, :im_info[bid, 4].astype("int32")] # id 0 is used for background classes, so -1 first labels = boxes_info[:, 4].astype("int32") - 1 pred_box = self.box_coder.decode(anchors, pred_offsets[bid]).detach() overlaps = layers.get_iou(boxes_info[:, :4], pred_box).detach() thresh1 = self.cfg.box_iou_threshold thresh2 = F.clip(overlaps.max(axis=1, keepdims=True), lower=thresh1 + clamp_eps, upper=1.0) gt_pred_prob = F.clip((overlaps - thresh1) / (thresh2 - thresh1), lower=0, upper=1.0) image_boxes_prob = F.zeros(pred_logits.shape[1:]).detach() # guarantee that nonzero_idx is not empty if gt_pred_prob.max() > clamp_eps: _, nonzero_idx = F.cond_take(gt_pred_prob != 0, gt_pred_prob) # since nonzeros is only 1 dim, use num_anchor to get real indices num_anchors = gt_pred_prob.shape[1] anchors_idx = nonzero_idx % num_anchors gt_idx = nonzero_idx // num_anchors image_boxes_prob[anchors_idx, labels[gt_idx]] = gt_pred_prob[gt_idx, anchors_idx] box_prob_list.append(image_boxes_prob) # construct bags for objects match_quality_matrix = layers.get_iou(boxes_info[:, :4], anchors).detach() num_gt = match_quality_matrix.shape[0] _, matched_idx = F.topk( match_quality_matrix, k=bucket_size, descending=True, no_sort=True, ) matched_idx = matched_idx.detach() matched_idx_flatten = matched_idx.reshape(-1) gather_idx = labels.reshape(-1, 1) gather_idx = F.broadcast_to(gather_idx, (num_gt, bucket_size)) gather_src = pred_scores[bid, matched_idx_flatten] gather_src = gather_src.reshape(num_gt, bucket_size, -1) matched_score = F.indexing_one_hot(gather_src, gather_idx, axis=2) topk_anchors = anchors[matched_idx_flatten] boxes_broad_cast = F.broadcast_to( F.expand_dims(boxes_info[:, :4], axis=1), (num_gt, bucket_size, 4)).reshape(-1, 4) matched_offsets = self.box_coder.encode(topk_anchors, boxes_broad_cast) reg_loss = layers.smooth_l1_loss( pred_offsets[bid, matched_idx_flatten], matched_offsets, beta=self.cfg.smooth_l1_beta).sum( axis=-1) * self.cfg.reg_loss_weight matched_reg_scores = F.exp(-reg_loss) positive_losses.append( positive_bag_loss(matched_score * matched_reg_scores.reshape(-1, bucket_size), axis=1)) num_foreground = im_info[:, 4].sum() pos_loss = F.concat(positive_losses).sum() / F.maximum( 1.0, num_foreground) box_probs = F.stack(box_prob_list, axis=0) neg_loss = negative_bag_loss( pred_scores * (1 - box_probs), self.cfg.focal_loss_gamma).sum() / F.maximum( 1.0, num_foreground * bucket_size) alpha = self.cfg.focal_loss_alpha pos_loss = pos_loss * alpha neg_loss = neg_loss * (1 - alpha) loss_dict = { "total_loss": pos_loss + neg_loss, "pos_loss": pos_loss, "neg_loss": neg_loss, } return loss_dict
def main(): parser = argparse.ArgumentParser() parser.add_argument("-a", "--arch", default="resnet18", type=str) parser.add_argument("-c", "--checkpoint", default=None, type=str) parser.add_argument("-i", "--image", default=None, type=str) parser.add_argument( "-m", "--mode", default="quantized", type=str, choices=["normal", "qat", "quantized"], help="Quantization Mode\n" "normal: no quantization, using float32\n" "qat: quantization aware training, simulate int8\n" "quantized: convert mode to int8 quantized, inference only", ) parser.add_argument("--dump", action="store_true", help="Dump quantized model") args = parser.parse_args() model = models.__dict__[args.arch]() if args.mode != "normal": quantize_qat(model, qconfig=Q.ema_fakequant_qconfig) if args.mode == "quantized": quantize(model) if args.checkpoint: logger.info("Load pretrained weights from %s", args.checkpoint) ckpt = mge.load(args.checkpoint) ckpt = ckpt["state_dict"] if "state_dict" in ckpt else ckpt model.load_state_dict(ckpt, strict=False) rpath = os.path.realpath(__file__ + "/../../") if args.image is None: path = rpath + "/assets/cat.jpg" else: path = args.image image = cv2.imread(path, cv2.IMREAD_COLOR) transform = T.Compose( [T.Resize(256), T.CenterCrop(224), T.Normalize(mean=128), T.ToMode("CHW")] ) @trace(symbolic=True, capture_as_const=True) def infer_func(processed_img): model.eval() logits = model(processed_img) probs = F.softmax(logits) return probs processed_img = transform.apply(image)[np.newaxis, :] processed_img = mge.tensor(processed_img, dtype="float32") probs = infer_func(processed_img) top_probs, classes = F.topk(probs, k=5, descending=True) if args.dump: output_file = ".".join([args.arch, args.mode, "megengine"]) logger.info("Dump to {}".format(output_file)) infer_func.dump(output_file, arg_names=["data"]) mge.save(model.state_dict(), output_file.replace("megengine", "pkl")) with open(rpath + "/assets/imagenet_class_info.json") as fp: imagenet_class_index = json.load(fp) for rank, (prob, classid) in enumerate( zip(top_probs.numpy().reshape(-1), classes.numpy().reshape(-1)) ): print( "{}: class = {:20s} with probability = {:4.1f} %".format( rank, imagenet_class_index[str(classid)][1], 100 * prob ) )