def test_pad(self):
# ***************************************************************
# Test ReplicationPad2d Layer
# ***************************************************************
arr = np.random.randn(16,3,224,224)
check_equal(arr, jnn.ReplicationPad2d(10), tnn.ReplicationPad2d(10))
check_equal(arr, jnn.ReplicationPad2d((1,23,4,5)), tnn.ReplicationPad2d((1,23,4,5)))
check_equal(arr, jnn.ReplicationPad2d((1,0,1,5)), tnn.ReplicationPad2d((1,0,1,5)))
check_equal(arr, jnn.ReplicationPad2d((100)), tnn.ReplicationPad2d((100)))
# ***************************************************************
# Test ConstantPad2d Layer
# ***************************************************************
arr = np.random.randn(16,3,224,224)
check_equal(arr, jnn.ConstantPad2d(10,-2), tnn.ConstantPad2d(10,-2))
check_equal(arr, jnn.ConstantPad2d((2,3,34,1),10.2), tnn.ConstantPad2d((2,3,34,1),10.2))
# ***************************************************************
# Test ZeroPad2d Layer
# ***************************************************************
arr = np.random.randn(16,3,224,224)
check_equal(arr, jnn.ZeroPad2d(1), tnn.ZeroPad2d(1))
check_equal(arr, jnn.ZeroPad2d((2,3,34,1)), tnn.ZeroPad2d((2,3,34,1)))
# ***************************************************************
# Test ReflectionPad2d Layer
# ***************************************************************
arr = np.random.randn(16,3,224,224)
check_equal(arr, jnn.ReflectionPad2d(20), tnn.ReflectionPad2d(20))
check_equal(arr, jnn.ReflectionPad2d((2,3,34,1)), tnn.ReflectionPad2d((2,3,34,1)))
check_equal(arr, jnn.ReflectionPad2d((10,123,34,1)), tnn.ReflectionPad2d((10,123,34,1)))
check_equal(arr, jnn.ReflectionPad2d((100)), tnn.ReflectionPad2d((100)))
def test_pad(self):
# ***************************************************************
# Test ReplicationPad2d Layer
# ***************************************************************
arr = np.random.randn(16, 3, 224, 224)
check_equal(arr, jnn.ReplicationPad2d(10), tnn.ReplicationPad2d(10))
check_equal(arr, jnn.ReplicationPad2d((1, 23, 4, 5)),
tnn.ReplicationPad2d((1, 23, 4, 5)))
check_equal(arr, jnn.ReplicationPad2d((1, 0, 1, 5)),
tnn.ReplicationPad2d((1, 0, 1, 5)))
check_equal(arr, jnn.ReplicationPad2d((100)),
tnn.ReplicationPad2d((100)))
# ***************************************************************
# Test ConstantPad2d Layer
# ***************************************************************
arr = np.random.randn(16, 3, 224, 224)
check_equal(arr, jnn.ConstantPad2d(10, -2), tnn.ConstantPad2d(10, -2))
check_equal(arr, jnn.ConstantPad2d((2, 3, 34, 1), 10.2),
tnn.ConstantPad2d((2, 3, 34, 1), 10.2))
arr = np.random.randn(16, 3, 224, 10, 10)
check_equal(arr, jnn.ConstantPad2d(10, -2), tnn.ConstantPad2d(10, -2))
check_equal(arr, jnn.ConstantPad2d((2, 3, 34, 1), 10.2),
tnn.ConstantPad2d((2, 3, 34, 1), 10.2))
# ***************************************************************
# Test ZeroPad2d Layer
# ***************************************************************
arr = np.random.randn(16, 3, 224, 224)
check_equal(arr, jnn.ZeroPad2d(1), tnn.ZeroPad2d(1))
check_equal(arr, jnn.ZeroPad2d((2, 3, 34, 1)),
tnn.ZeroPad2d((2, 3, 34, 1)))
# ***************************************************************
# Test ReflectionPad2d Layer
# ***************************************************************
arr = np.random.randn(16, 3, 224, 224)
check_equal(arr, jnn.ReflectionPad2d(20), tnn.ReflectionPad2d(20))
check_equal(arr, jnn.ReflectionPad2d((2, 3, 34, 1)),
tnn.ReflectionPad2d((2, 3, 34, 1)))
check_equal(arr, jnn.ReflectionPad2d((10, 123, 34, 1)),
tnn.ReflectionPad2d((10, 123, 34, 1)))
check_equal(arr, jnn.ReflectionPad2d((100)), tnn.ReflectionPad2d(
(100)))
# ***************************************************************
# Test function pad
# ***************************************************************
arr = np.random.randn(16, 3, 224, 224)
padding = (10, 11, 2, 3)
for mode in ['constant', 'replicate', 'reflect', 'circular']:
j_data = jt.array(arr)
t_data = torch.tensor(arr)
t_output = tnn.functional.pad(t_data, padding,
mode=mode).detach().numpy()
j_output = jnn.pad(j_data, padding, mode).numpy()
assert np.allclose(t_output, j_output)
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