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