def execute(self, x):
dims = [0] + list(range(2, x.ndim))
####### centering calibration begin ####### $
x += self.center_weight * self.stas(x)
####### centering calibration end ####### $
if self.is_train:
xmean = jt.mean(x, dims=dims)
x2mean = jt.mean(x * x, dims=dims)
if self.sync and jt.in_mpi:
xmean = xmean.mpi_all_reduce("mean")
x2mean = x2mean.mpi_all_reduce("mean")
xvar = (x2mean - xmean * xmean).maximum(0.0)
w = 1.0 / jt.sqrt(xvar + self.eps)
b = -xmean * w
norm_x = x * w.broadcast(x, dims) + b.broadcast(x, dims)
self.running_mean.update(self.running_mean + (xmean.reshape(
(-1, )) - self.running_mean) * self.momentum)
self.running_var.update(self.running_var +
(xvar.reshape((-1, )) - self.running_var) *
self.momentum)
else:
w = 1.0 / jt.sqrt(self.running_var + self.eps)
b = -self.running_mean * w
norm_x = x * w.broadcast(x, dims) + b.broadcast(x, dims)
####### scaling calibration begin ####### $
scale_factor = jt.sigmoid(self.scale_weight * self.stas(norm_x) +
self.scale_bias)
####### scaling calibration end ####### $
return self.weight * scale_factor * norm_x + self.bias
def execute(self, pred, true):
loss = self.loss_fcn(pred, true)
pred = jt.sigmoid(pred) # prob from logits
dx = pred - true # reduce only missing label effects
# dx = (pred - true).abs() # reduce missing label and false label effects
alpha_factor = 1 - jt.exp((dx - 1) / (self.alpha + 1e-4))
loss *= alpha_factor
return loss.mean()
def __init__(self, p=None, logits=None):
assert (p is not None) or (logits is not None)
assert 0 < p and p < 1
if p is None:
self.prob = jt.sigmoid(logits)
self.logits = logits
elif logits is None:
self.prob = p
self.logits = -jt.safe_log(1. / p - 1)
def execute(self, x):
y = x[0] # no weight
if self.weight:
w = jt.sigmoid(self.w) * 2
for i in self.iter:
y = y + x[i + 1] * w[i]
else:
for i in self.iter:
y = y + x[i + 1]
return y
def __init__(self, probs=None, logits=None):
assert not (probs is None and logits is None)
if probs is None:
# cannot align to pytorch
probs = jt.sigmoid(logits)
with jt.no_grad():
self.probs = probs / probs.sum(-1, True)
self.cum_probs = simple_presum(probs)
self.cum_probs_l = self.cum_probs[..., :-1]
self.cum_probs_r = self.cum_probs[..., 1:]
def execute(self, x):
x = self.resnet.conv1(x)
x = self.resnet.bn1(x)
x = nn.relu(x)
x = self.resnet.maxpool(x)
x1 = self.resnet.layer1(x)
x2 = self.resnet.layer2(x1)
x3 = self.resnet.layer3(x2)
x4 = self.resnet.layer4(x3)
x2_rfb = self.rfb2_1(x2)
x3_rfb = self.rfb3_1(x3)
x4_rfb = self.rfb4_1(x4)
ra5_feat = self.agg1(x4_rfb, x3_rfb, x2_rfb)
lateral_map_5 = self.upsample1_1(ra5_feat)
crop_4 = self.upsample1_2(ra5_feat)
x = (((-1) * jt.sigmoid(crop_4)) + 1)
x = x.expand((-1), 2048, (-1), (-1)).multiply(x4)
x = self.ra4_conv1(x)
x = nn.relu(self.ra4_conv2(x))
x = nn.relu(self.ra4_conv3(x))
x = nn.relu(self.ra4_conv4(x))
ra4_feat = self.ra4_conv5(x)
x = (ra4_feat + crop_4)
lateral_map_4 = self.upsample2_1(x)
crop_3 = self.upsample2_2(x)
x = (((-1) * jt.sigmoid(crop_3)) + 1)
x = x.expand((-1), 1024, (-1), (-1)).multiply(x3)
x = self.ra3_conv1(x)
x = nn.relu(self.ra3_conv2(x))
x = nn.relu(self.ra3_conv3(x))
ra3_feat = self.ra3_conv4(x)
x = (ra3_feat + crop_3)
lateral_map_3 = self.upsample3_1(x)
crop_2 = self.upsample3_2(x)
x = (((-1) * jt.sigmoid(crop_2)) + 1)
x = x.expand((-1), 512, (-1), (-1)).multiply(x2)
x = self.ra2_conv1(x)
x = nn.relu(self.ra2_conv2(x))
x = nn.relu(self.ra2_conv3(x))
ra2_feat = self.ra2_conv4(x)
x = (ra2_feat + crop_2)
lateral_map_2 = self.upsample4(x)
return (lateral_map_5, lateral_map_4, lateral_map_3, lateral_map_2)
def execute(self, pred, true):
loss = self.loss_fcn(pred, true)
pred_prob = jt.sigmoid(pred) # prob from logits
alpha_factor = true * self.alpha + (1 - true) * (1 - self.alpha)
modulating_factor = jt.abs(true - pred_prob)**self.gamma
loss *= alpha_factor * modulating_factor
if self.reduction == 'mean':
return loss.mean()
elif self.reduction == 'sum':
return loss.sum()
else: # 'none'
return loss
def plot_wh_methods(): # from utils.plots import *; plot_wh_methods()
# Compares the two methods for width-height anchor multiplication
# https://github.com/ultralytics/yolov3/issues/168
x = np.arange(-4.0, 4.0, .1)
ya = np.exp(x)
yb = jt.sigmoid(jt.array(x)).numpy() * 2
fig = plt.figure(figsize=(6, 3), tight_layout=True)
plt.plot(x, ya, '.-', label='YOLOv3')
plt.plot(x, yb ** 2, '.-', label='YOLOv5 ^2')
plt.plot(x, yb ** 1.6, '.-', label='YOLOv5 ^1.6')
plt.xlim(left=-4, right=4)
plt.ylim(bottom=0, top=6)
plt.xlabel('input')
plt.ylabel('output')
plt.grid()
plt.legend()
fig.savefig('comparison.png', dpi=200)
def execute(self, pred, true):
loss = self.loss_fcn(pred, true)
# p_t = jt.exp(-loss)
# loss *= self.alpha * (1.000001 - p_t) ** self.gamma # non-zero power for gradient stability
# TF implementation https://github.com/tensorflow/addons/blob/v0.7.1/tensorflow_addons/losses/focal_loss.py
pred_prob = jt.sigmoid(pred) # prob from logits
p_t = true * pred_prob + (1 - true) * (1 - pred_prob)
alpha_factor = true * self.alpha + (1 - true) * (1 - self.alpha)
modulating_factor = (1.0 - p_t)**self.gamma
loss *= alpha_factor * modulating_factor
if self.reduction == 'mean':
return loss.mean()
elif self.reduction == 'sum':
return loss.sum()
else: # 'none'
return loss
def integrator(raw, z_vals, rays_d, raw_noise_std=0, white_bkgd=False):
"""Transforms model's predictions to semantically meaningful values.
Args:
raw: [num_rays, num_samples along ray, 4]. Prediction from model.
z_vals: [num_rays, num_samples along ray]. Integration time.
rays_d: [num_rays, 3]. Direction of each ray.
Returns:
rgb_map: [num_rays, 3]. Estimated RGB color of a ray.
disp_map: [num_rays]. Disparity map. Inverse of depth map.
acc_map: [num_rays]. Sum of weights along each ray.
weights: [num_rays, num_samples]. Weights assigned to each sampled color.
depth_map: [num_rays]. Estimated distance to object.
"""
raw2alpha = lambda raw, dists, act_fn=jt.nn.relu: 1. - jt.exp(-act_fn(raw)
* dists)
dists = z_vals[..., 1:] - z_vals[..., :-1]
dists = jt.concat([
dists,
jt.array(np.array([1e10]).astype(np.float32)).expand(
dists[..., :1].shape)
], -1) # [N_rays, N_samples]
dists = dists * jt.norm(rays_d.unsqueeze(-2), p=2, dim=-1)
rgb = jt.sigmoid(raw[..., :3]) # [N_rays, N_samples, 3]
noise = 0.
if raw_noise_std > 0.:
noise = jt.init.gauss(raw[..., 3].shape, raw.dtype) * raw_noise_std
alpha = raw2alpha(raw[..., 3] + noise, dists) # [N_rays, N_samples]
weights = alpha * jt.cumprod(
jt.concat([jt.ones(
(alpha.shape[0], 1)), 1. - alpha + 1e-10], -1), -1)[:, :-1]
rgb_map = jt.sum(weights.unsqueeze(-1) * rgb, -2) # [N_rays, 3]
depth_map = jt.sum(weights * z_vals, -1)
disp_map = 1. / jt.maximum(1e-10 * jt.ones_like(depth_map),
depth_map / jt.sum(weights, -1))
acc_map = jt.sum(weights, -1)
if white_bkgd:
rgb_map = rgb_map + (1. - acc_map.unsqueeze(-1))
return rgb_map, disp_map, acc_map, weights, depth_map
def log_sigmoid(x):
return jt.log(jt.sigmoid(x))
def execute(self, x):
""" The input should be of size [batch_size, 3, img_h, img_w] """
_, _, img_h, img_w = x.shape
cfg._tmp_img_h = img_h
cfg._tmp_img_w = img_w
with timer.env('backbone'):
outs = self.backbone(x)
if cfg.fpn is not None:
with timer.env('fpn'):
# Use backbone.selected_layers because we overwrote self.selected_layers
outs = [outs[i] for i in cfg.backbone.selected_layers]
outs = self.fpn(outs)
proto_out = None
if cfg.mask_type == mask_type.lincomb and cfg.eval_mask_branch:
with timer.env('proto'):
proto_x = x if self.proto_src is None else outs[self.proto_src]
if self.num_grids > 0:
grids = self.grid.repeat(proto_x.shape[0], 1, 1, 1)
proto_x = jt.contrib.concat([proto_x, grids], dim=1)
proto_out = self.proto_net(proto_x)
proto_out = cfg.mask_proto_prototype_activation(proto_out)
if cfg.mask_proto_prototypes_as_features:
# Clone here because we don't want to permute this, though idk if contiguous makes this unnecessary
proto_downsampled = proto_out.clone()
if cfg.mask_proto_prototypes_as_features_no_grad:
proto_downsampled = proto_out.detach()
# Move the features last so the multiplication is easy
proto_out = proto_out.permute(0, 2, 3, 1)
if cfg.mask_proto_bias:
bias_shape = [x for x in proto_out.shape]
bias_shape[-1] = 1
proto_out = jt.contrib.concat(
[proto_out, jt.ones(bias_shape)], -1)
with timer.env('pred_heads'):
pred_outs = {'loc': [], 'conf': [], 'mask': [], 'priors': []}
if cfg.use_mask_scoring:
pred_outs['score'] = []
if cfg.use_instance_coeff:
pred_outs['inst'] = []
for idx, pred_layer in zip(self.selected_layers,
self.prediction_layers):
pred_x = outs[idx]
if cfg.mask_type == mask_type.lincomb and cfg.mask_proto_prototypes_as_features:
# Scale the prototypes down to the current prediction layer's size and add it as inputs
proto_downsampled = nn.interpolate(
proto_downsampled,
size=outs[idx].shape[2:],
mode='bilinear',
align_corners=False)
# proto_downsampled = interpolate(proto_downsampled, size=outs[idx].shape[2:], mode='bilinear', align_corners=False)
pred_x = jt.contrib.concat([pred_x, proto_downsampled],
dim=1)
# A hack for the way dataparallel works
if cfg.share_prediction_module and pred_layer is not self.prediction_layers[
0]:
pred_layer.parent = [self.prediction_layers[0]]
p = pred_layer(pred_x)
for k, v in p.items():
pred_outs[k].append(v)
for k, v in pred_outs.items():
pred_outs[k] = jt.contrib.concat(v, -2)
if proto_out is not None:
pred_outs['proto'] = proto_out
#print('hh',pred_outs)
#print()
if self.is_training():
# For the extra loss functions
if cfg.use_class_existence_loss:
pred_outs['classes'] = self.class_existence_fc(
outs[-1].mean(dim=(2, 3)))
if cfg.use_semantic_segmentation_loss:
pred_outs['segm'] = self.semantic_seg_conv(outs[0])
return pred_outs
else:
if cfg.use_mask_scoring:
pred_outs['score'] = jt.sigmoid(pred_outs['score'])
if cfg.use_focal_loss:
if cfg.use_sigmoid_focal_loss:
# Note: even though conf[0] exists, this mode doesn't train it so don't use it
pred_outs['conf'] = jt.sigmoid(pred_outs['conf'])
if cfg.use_mask_scoring:
pred_outs['conf'] *= pred_outs['score']
elif cfg.use_objectness_score:
# See focal_loss_sigmoid in multibox_loss.py for details
objectness = jt.sigmoid(pred_outs['conf'][:, :, 0])
pred_outs['conf'][:, :, 1:] = objectness.unsqueeze(
2) * nn.softmax(pred_outs['conf'][:, :, 1:], -1)
pred_outs['conf'][:, :, 0] = 1 - objectness
else:
pred_outs['conf'] = nn.softmax(pred_outs['conf'], -1)
else:
if cfg.use_objectness_score:
objectness = jt.sigmoid(pred_outs['conf'][:, :, 0])
pred_outs['conf'][:, :, 1:] = (objectness > 0.10).unsqueeze(-1) \
* nn.softmax(pred_outs['conf'][:, :, 1:], dim=-1)
else:
pred_outs['conf'] = nn.softmax(pred_outs['conf'], -1)
return self.detect(pred_outs, self)
def execute(self, x):
"""
Args:
- x: The convOut from a layer in the backbone network
Size: [batch_size, in_channels, conv_h, conv_w])
Returns a tuple (bbox_coords, class_confs, mask_output, prior_boxes) with sizes
- bbox_coords: [batch_size, conv_h*conv_w*num_priors, 4]
- class_confs: [batch_size, conv_h*conv_w*num_priors, num_classes]
- mask_output: [batch_size, conv_h*conv_w*num_priors, mask_dim]
- prior_boxes: [conv_h*conv_w*num_priors, 4]
"""
# In case we want to use another module's layers
src = self if self.parent[0] is None else self.parent[0]
conv_h = x.shape[2]
conv_w = x.shape[3]
if cfg.extra_head_net is not None:
x = src.upfeature(x)
if cfg.use_prediction_module:
# The two branches of PM design (c)
a = src.block(x)
b = src.conv(x)
b = src.bn(b)
b = nn.relu(b)
# TODO: Possibly switch this out for a product
x = a + b
bbox_x = src.bbox_extra(x)
conf_x = src.conf_extra(x)
mask_x = src.mask_extra(x)
bbox = src.bbox_layer(bbox_x).permute(0, 2, 3,
1).view(x.shape[0], -1, 4)
conf = src.conf_layer(conf_x).permute(0, 2, 3,
1).view(x.shape[0], -1,
self.num_classes)
if cfg.eval_mask_branch:
mask = src.mask_layer(mask_x).permute(0, 2, 3, 1).view(
x.shape[0], -1, self.mask_dim)
else:
mask = jt.zeros((x.shape[0], bbox.shape[1], self.mask_dim))
if cfg.use_mask_scoring:
score = src.score_layer(x).permute(0, 2, 3,
1).view(x.shape[0], -1, 1)
if cfg.use_instance_coeff:
inst = src.inst_layer(x).permute(0, 2, 3,
1).view(x.shape[0], -1,
cfg.num_instance_coeffs)
# See box_utils.decode for an explanation of this
if cfg.use_yolo_regressors:
bbox[:, :, :2] = jt.sigmoid(bbox[:, :, :2]) - 0.5
bbox[:, :, 0] /= conv_w
bbox[:, :, 1] /= conv_h
if cfg.eval_mask_branch:
if cfg.mask_type == mask_type.direct:
mask = jt.sigmoid(mask)
elif cfg.mask_type == mask_type.lincomb:
mask = cfg.mask_proto_coeff_activation(mask)
if cfg.mask_proto_coeff_gate:
gate = src.gate_layer(x).permute(0, 2, 3, 1).view(
x.shape[0], -1, self.mask_dim)
mask = mask * jt.sigmoid(gate)
if cfg.mask_proto_split_prototypes_by_head and cfg.mask_type == mask_type.lincomb:
mask = nn.ConstantPad2d(
(self.index * self.mask_dim,
(self.num_heads - self.index - 1) * self.mask_dim),
value=0)(mask)
priors = self.make_priors(conv_h, conv_w)
preds = {'loc': bbox, 'conf': conf, 'mask': mask, 'priors': priors}
if cfg.use_mask_scoring:
preds['score'] = score
if cfg.use_instance_coeff:
preds['inst'] = inst
return preds
def execute(x):
return x * jt.sigmoid(x)
def grad(self, grad_output):
x = self.x
sx = jt.sigmoid(x)
fx = jt.softplus(x).tanh()
return grad_output * (fx + x * sx * (1 - fx * fx))
def grad(self, grad_output):
x = self.x
sx = jt.sigmoid(x)
return grad_output * (sx * (1 + x * (1 - sx)))
def execute(self, x):
self.x = x
return x * jt.sigmoid(x)
