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
