def test_zeros():
assertTensorClose(
mge.zeros((2, 2), dtype=np.int32).numpy(),
np.zeros((2, 2), dtype=np.int32))
assertTensorClose(
mge.zeros(mge.tensor([2, 2], dtype=np.int32), dtype=np.int32).numpy(),
np.zeros((2, 2), dtype=np.int32),
)
def roi_pool(rpn_fms, rois, stride, pool_shape, roi_type='roi_align',
labels=None, bbox_targets=None):
assert len(stride) == len(rpn_fms)
canonical_level = 4
canonical_box_size = 224
min_level = math.log2(stride[0])
max_level = math.log2(stride[-1])
num_fms = len(rpn_fms)
box_sizes = F.sqrt((rois[:, 3] - rois[:, 1]) * (rois[:, 4] - rois[:, 2]))
level_assignments = F.floor(
canonical_level + F.log(box_sizes / canonical_box_size) / np.log(2)
)
level_assignments = F.minimum(level_assignments, max_level)
level_assignments = F.maximum(level_assignments, min_level)
level_assignments = level_assignments - min_level
available_masks = F.concat(
[mge.ones(level_assignments.shapeof()[0]), mge.zeros(num_fms)], axis=0)
level_assignments = F.concat([level_assignments, mge.tensor(np.arange(num_fms, dtype=np.int32))], axis=0)
rois = F.concat([rois, mge.zeros((num_fms, rois.shapeof()[-1]))], axis=0)
if labels is not None:
labels = F.concat([labels, mge.ones((num_fms, labels.shapeof()[-1]))], axis=0)
bbox_targets = F.concat([bbox_targets, mge.zeros((num_fms, bbox_targets.shapeof()[-1]))], axis=0)
pool_list, inds_list = [], []
for i in range(len(rpn_fms)):
mask = level_assignments == i
inds = mask_to_inds(mask)
rois_fm = rois.ai[inds]
if roi_type == 'roi_pool':
pool_fm = F.roi_pooling(
rpn_fms[i], rois_fm, pool_shape, mode='max', scale=1.0/stride[i])
elif roi_type == 'roi_align':
pool_fm = F.roi_align(
rpn_fms[i], rois_fm, pool_shape, mode='average',
spatial_scale=1.0/stride[i], sample_points=2, aligned=True)
pool_list.append(pool_fm)
inds_list.append(inds)
fm_order = F.concat(inds_list, axis=0)
pool_feature = F.concat(pool_list, axis=0)
ordered_available_masks = available_masks.ai[fm_order]
available_inds = mask_to_inds(ordered_available_masks)
pool_feature = pool_feature.ai[available_inds]
rois = rois.ai[fm_order, :].ai[available_inds, :]
if labels is not None:
labels = labels.ai[fm_order].ai[available_inds]
bbox_targets = bbox_targets.ai[fm_order, :].ai[available_inds, :]
return pool_feature, rois, F.zero_grad(labels), F.zero_grad(bbox_targets)
else:
return pool_feature, rois, None, None
def roi_pool(
rpn_fms,
rois,
stride,
pool_shape,
roi_type="roi_align",
):
assert len(stride) == len(rpn_fms)
canonical_level = 4
canonical_box_size = 224
min_level = math.log2(stride[0])
max_level = math.log2(stride[-1])
num_fms = len(rpn_fms)
box_area = (rois[:, 3] - rois[:, 1]) * (rois[:, 4] - rois[:, 2])
level_assignments = F.floor(canonical_level +
F.log(box_area.sqrt() / canonical_box_size) /
np.log(2))
level_assignments = F.minimum(level_assignments, max_level)
level_assignments = F.maximum(level_assignments, min_level)
level_assignments = level_assignments - min_level
# avoid empty assignment
level_assignments = F.concat(
[level_assignments,
mge.tensor(np.arange(num_fms, dtype=np.int32))], )
rois = F.concat([rois, mge.zeros((num_fms, rois.shapeof(-1)))])
pool_list, inds_list = [], []
for i in range(num_fms):
mask = level_assignments == i
_, inds = F.cond_take(mask == 1, mask)
level_rois = rois.ai[inds]
if roi_type == "roi_pool":
pool_fm = F.roi_pooling(rpn_fms[i],
level_rois,
pool_shape,
mode="max",
scale=1.0 / stride[i])
elif roi_type == "roi_align":
pool_fm = F.roi_align(
rpn_fms[i],
level_rois,
pool_shape,
mode="average",
spatial_scale=1.0 / stride[i],
sample_points=2,
aligned=True,
)
pool_list.append(pool_fm)
inds_list.append(inds)
fm_order = F.concat(inds_list, axis=0)
fm_order = F.argsort(fm_order.reshape(1, -1))[1].reshape(-1)
pool_feature = F.concat(pool_list, axis=0)
pool_feature = pool_feature.ai[fm_order][:-num_fms]
return pool_feature
def add_H_W_Padding(x, margin):
shape = x.shape
padding_shape = list(shape)[:-2] + [
shape[-2] + 2 * margin, shape[-1] + 2 * margin
]
res = mge.zeros(padding_shape, dtype=x.dtype)
res = res.set_subtensor(x)[:, :, margin:margin + shape[-2],
margin:margin + shape[-1]]
return res
def get_focal_loss(
score: Tensor,
label: Tensor,
ignore_label: int = -1,
background: int = 0,
alpha: float = 0.5,
gamma: float = 0,
norm_type: str = "fg",
) -> Tensor:
r"""Focal Loss for Dense Object Detection:
<https://arxiv.org/pdf/1708.02002.pdf>
.. math::
FL(p_t) = -\alpha_t(1-p_t)^\gamma \log(p_t)
Args:
score (Tensor):
the predicted score with the shape of :math:`(B, A, C)`
label (Tensor):
the assigned label of boxes with shape of :math:`(B, A)`
ignore_label (int):
the value of ignore class. Default: -1
background (int):
the value of background class. Default: 0
alpha (float):
parameter to mitigate class imbalance. Default: 0.5
gamma (float):
parameter to mitigate easy/hard loss imbalance. Default: 0
norm_type (str): current support 'fg', 'none':
'fg': loss will be normalized by number of fore-ground samples
'none": not norm
Returns:
the calculated focal loss.
"""
mask = 1 - (label == ignore_label)
valid_label = label * mask
score_shp = score.shape
zero_mat = mge.zeros(
F.concat([score_shp[0], score_shp[1], score_shp[2] + 1], axis=0),
dtype=np.float32,
)
one_mat = mge.ones(
F.concat([score_shp[0], score_shp[1], tensor(1)], axis=0), dtype=np.float32,
)
one_hot = basic.indexing_set_one_hot(
zero_mat, 2, valid_label.astype(np.int32), one_mat
)[:, :, 1:]
pos_part = F.power(1 - score, gamma) * one_hot * F.log(score)
neg_part = F.power(score, gamma) * (1 - one_hot) * F.log(1 - score)
loss = -(alpha * pos_part + (1 - alpha) * neg_part).sum(axis=2) * mask
if norm_type == "fg":
positive_mask = label > background
return loss.sum() / F.maximum(positive_mask.sum(), 1)
elif norm_type == "none":
return loss.sum()
else:
raise NotImplementedError
def net_stats(model, input_size, bar_length_max=20, log_params=True, log_flops=True):
def dict2table(list_of_dict, header):
table_data = [header]
for d in list_of_dict:
row = []
for h in header:
v = ""
if h in d:
v = d[h]
row.append(v)
table_data.append(row)
return table_data
def sizeof_fmt(num, suffix="B"):
for unit in ["", "Ki", "Mi", "Gi", "Ti", "Pi", "Ei", "Zi"]:
if abs(num) < 1024.0:
return "{:3.3f} {}{}".format(num, unit, suffix)
num /= 1024.0
sign_str = "-" if num < 0 else ""
return "{}{:.1f} {}{}".format(sign_str, num, "Yi", suffix)
def get_byteswidth(tensor):
dtype = tensor.dtype
if mgb.dtype.is_quantize(dtype):
return 1
elif mgb.dtype.is_bfloat16(dtype):
return 2
else:
return 4
def print_flops_stats(flops):
flops_list = [i["flops_num"] for i in flops]
max_flops_num = max(flops_list + [0])
# calc total flops and set flops_cum
total_flops_num = 0
for d in flops:
total_flops_num += int(d["flops_num"])
d["flops_cum"] = sizeof_fmt(total_flops_num, suffix="OPs")
for i in flops:
f = i["flops_num"]
i["flops"] = sizeof_fmt(f, suffix="OPs")
r = i["ratio"] = f / total_flops_num
i["percentage"] = "{:.2f}%".format(r * 100)
bar_length = int(f / max_flops_num * bar_length_max)
i["bar"] = "#" * bar_length
header = [
"name",
"class_name",
"input_shapes",
"output_shapes",
"flops",
"flops_cum",
"percentage",
"bar",
]
total_flops_str = sizeof_fmt(total_flops_num, suffix="OPs")
total_var_size = sum(sum(s[1] for s in i["output_shapes"]) for i in flops)
flops.append(
dict(name="total", flops=total_flops_str, output_shapes=total_var_size)
)
logger.info(
"flops stats: \n" + tabulate.tabulate(dict2table(flops, header=header))
)
return total_flops_num
def print_params_stats(params):
total_param_dims, total_param_size = 0, 0
for d in params:
total_param_dims += int(d["param_dim"])
total_param_size += int(d["size"])
d["size"] = sizeof_fmt(d["size"])
d["size_cum"] = sizeof_fmt(total_param_size)
for d in params:
ratio = d["param_dim"] / total_param_dims
d["ratio"] = ratio
d["percentage"] = "{:.2f}%".format(ratio * 100)
# construct bar
max_ratio = max([d["ratio"] for d in params])
for d in params:
bar_length = int(d["ratio"] / max_ratio * bar_length_max)
d["size_bar"] = "#" * bar_length
param_size = sizeof_fmt(total_param_size)
params.append(dict(name="total", param_dim=total_param_dims, size=param_size,))
header = [
"name",
"shape",
"mean",
"std",
"param_dim",
"bits",
"size",
"size_cum",
"percentage",
"size_bar",
]
logger.info(
"param stats: \n" + tabulate.tabulate(dict2table(params, header=header))
)
return total_param_size
def net_stats_hook(module, input, output, name=""):
class_name = str(module.__class__).split(".")[-1].split("'")[0]
flops_fun = CALC_FLOPS.get(type(module))
if callable(flops_fun):
flops_num = flops_fun(module, input, output)
if not isinstance(output, (list, tuple)):
output = [output]
flops.append(
dict(
name=name,
class_name=class_name,
input_shapes=[i.shape for i in input],
output_shapes=[o.shape for o in output],
flops_num=flops_num,
flops_cum=0,
)
)
if hasattr(module, "weight") and module.weight is not None:
w = module.weight
value = w.numpy()
param_dim = np.prod(w.shape)
param_bytes = get_byteswidth(w)
params.append(
dict(
name=name + "-w",
shape=w.shape,
param_dim=param_dim,
bits=param_bytes * 8,
size=param_dim * param_bytes,
size_cum=0,
mean="{:.2g}".format(value.mean()),
std="{:.2g}".format(value.std()),
)
)
if hasattr(module, "bias") and module.bias is not None:
b = module.bias
value = b.numpy()
param_dim = np.prod(b.shape)
param_bytes = get_byteswidth(b)
params.append(
dict(
name=name + "-b",
shape=b.shape,
param_dim=param_dim,
bits=param_bytes * 8,
size=param_dim * param_bytes,
size_cum=0,
mean="{:.2g}".format(value.mean()),
std="{:.2g}".format(value.std()),
)
)
# multiple inputs to the network
if not isinstance(input_size[0], tuple):
input_size = [input_size]
params = []
flops = []
hooks = []
for (name, module) in model.named_modules():
if isinstance(module, hook_modules):
hooks.append(
module.register_forward_hook(partial(net_stats_hook, name=name))
)
inputs = [mge.zeros(in_size, dtype=np.float32) for in_size in input_size]
model.eval()
model(*inputs)
for h in hooks:
h.remove()
total_flops, total_params = 0, 0
if log_params:
total_params = print_params_stats(params)
if log_flops:
total_flops = print_flops_stats(flops)
return total_params, total_flops
def zeros_like(inp):
return mge.zeros(inp.shape, dtype=inp.dtype)