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
)
)
