def _init_layers(self): """Initialize layers of the head.""" super(FCOSHead, self)._init_cls_convs() super(FCOSHead, self)._init_reg_convs() self.relu = nn.ReLU(inplace=True) self.vfnet_reg_conv = ConvModule( self.feat_channels, self.feat_channels, 3, stride=1, padding=1, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg, bias=self.conv_bias) self.vfnet_reg = nn.Conv2d(self.feat_channels, 4, 3, padding=1) self.scales = nn.ModuleList([Scale(1.0) for _ in self.strides]) self.vfnet_reg_refine_dconv = DeformConv2d( self.feat_channels, self.feat_channels, self.dcn_kernel, 1, padding=self.dcn_pad) self.vfnet_reg_refine = nn.Conv2d(self.feat_channels, 4, 3, padding=1) self.scales_refine = nn.ModuleList([Scale(1.0) for _ in self.strides]) self.vfnet_cls_dconv = DeformConv2d( self.feat_channels, self.feat_channels, self.dcn_kernel, 1, padding=self.dcn_pad) self.vfnet_cls = nn.Conv2d( self.feat_channels, self.cls_out_channels, 3, padding=1)
def _init_centripetal_layers(self): """Initialize centripetal layers. Including feature adaption deform convs (feat_adaption), deform offset prediction convs (dcn_off), guiding shift (guiding_shift) and centripetal shift ( centripetal_shift). Each branch has two parts: prefix `tl_` for top-left and `br_` for bottom-right. """ self.tl_feat_adaption = nn.ModuleList() self.br_feat_adaption = nn.ModuleList() self.tl_dcn_offset = nn.ModuleList() self.br_dcn_offset = nn.ModuleList() self.tl_guiding_shift = nn.ModuleList() self.br_guiding_shift = nn.ModuleList() self.tl_centripetal_shift = nn.ModuleList() self.br_centripetal_shift = nn.ModuleList() for _ in range(self.num_feat_levels): self.tl_feat_adaption.append( DeformConv2d(self.in_channels, self.in_channels, self.feat_adaption_conv_kernel, 1, 1)) self.br_feat_adaption.append( DeformConv2d(self.in_channels, self.in_channels, self.feat_adaption_conv_kernel, 1, 1)) self.tl_guiding_shift.append( self._make_layers( out_channels=self.guiding_shift_channels, in_channels=self.in_channels)) self.br_guiding_shift.append( self._make_layers( out_channels=self.guiding_shift_channels, in_channels=self.in_channels)) self.tl_dcn_offset.append( ConvModule( self.guiding_shift_channels, self.feat_adaption_conv_kernel**2 * self.guiding_shift_channels, 1, bias=False, act_cfg=None)) self.br_dcn_offset.append( ConvModule( self.guiding_shift_channels, self.feat_adaption_conv_kernel**2 * self.guiding_shift_channels, 1, bias=False, act_cfg=None)) self.tl_centripetal_shift.append( self._make_layers( out_channels=self.centripetal_shift_channels, in_channels=self.in_channels)) self.br_centripetal_shift.append( self._make_layers( out_channels=self.centripetal_shift_channels, in_channels=self.in_channels))
def _test_amp_deformconv(self, input_dtype, threshold=1e-3): """The function to test amp released on pytorch 1.6.0. The type of input data might be torch.float or torch.half, so we should test deform_conv in both cases. With amp, the data type of model will NOT be set manually. Args: input_dtype: torch.float or torch.half. threshold: the same as above function. """ if not torch.cuda.is_available(): return from mmcv.ops import DeformConv2dPack c_in = 1 c_out = 1 x = torch.Tensor(input).cuda().type(input_dtype) x.requires_grad = True model = DeformConv2dPack(c_in, c_out, 2, stride=1, padding=0) model.conv_offset.weight.data = torch.nn.Parameter( torch.Tensor(offset_weight).reshape(8, 1, 2, 2)) model.conv_offset.bias.data = torch.nn.Parameter( torch.Tensor(offset_bias).reshape(8)) model.weight.data = torch.nn.Parameter( torch.Tensor(deform_weight).reshape(1, 1, 2, 2)) model.cuda() out = model(x) out.backward(torch.ones_like(out)) assert np.allclose(out.data.detach().cpu().numpy(), gt_out, threshold) assert np.allclose(x.grad.detach().cpu().numpy(), gt_x_grad, threshold) assert np.allclose( model.conv_offset.weight.grad.detach().cpu().numpy(), gt_offset_weight_grad, threshold) assert np.allclose(model.conv_offset.bias.grad.detach().cpu().numpy(), gt_offset_bias_grad, threshold) assert np.allclose(model.weight.grad.detach().cpu().numpy(), gt_deform_weight_grad, threshold) from mmcv.ops import DeformConv2d # test bias model = DeformConv2d(1, 1, 2, stride=1, padding=0) assert not hasattr(model, 'bias') # test bias=True with pytest.raises(AssertionError): model = DeformConv2d(1, 1, 2, stride=1, padding=0, bias=True) # test in_channels % group != 0 with pytest.raises(AssertionError): model = DeformConv2d(3, 2, 3, groups=2) # test out_channels % group != 0 with pytest.raises(AssertionError): model = DeformConv2d(3, 4, 3, groups=3)
def _test_deformconv(self, dtype=torch.float, threshold=1e-3, device='cuda'): if not torch.cuda.is_available() and device == 'cuda': pytest.skip('test requires GPU') from mmcv.ops import DeformConv2dPack c_in = 1 c_out = 1 x = torch.tensor(input, device=device, dtype=dtype) x.requires_grad = True model = DeformConv2dPack(c_in, c_out, 2, stride=1, padding=0) model.conv_offset.weight.data = torch.nn.Parameter( torch.Tensor(offset_weight).reshape(8, 1, 2, 2)) model.conv_offset.bias.data = torch.nn.Parameter( torch.Tensor(offset_bias).reshape(8)) model.weight.data = torch.nn.Parameter( torch.Tensor(deform_weight).reshape(1, 1, 2, 2)) if device == 'cuda': model.cuda() model.type(dtype) out = model(x) out.backward(torch.ones_like(out)) assert np.allclose(out.data.detach().cpu().numpy(), gt_out, threshold) assert np.allclose(x.grad.detach().cpu().numpy(), gt_x_grad, threshold) assert np.allclose( model.conv_offset.weight.grad.detach().cpu().numpy(), gt_offset_weight_grad, threshold) assert np.allclose(model.conv_offset.bias.grad.detach().cpu().numpy(), gt_offset_bias_grad, threshold) assert np.allclose(model.weight.grad.detach().cpu().numpy(), gt_deform_weight_grad, threshold) from mmcv.ops import DeformConv2d # test bias model = DeformConv2d(1, 1, 2, stride=1, padding=0) assert not hasattr(model, 'bias') # test bias=True with pytest.raises(AssertionError): model = DeformConv2d(1, 1, 2, stride=1, padding=0, bias=True) # test in_channels % group != 0 with pytest.raises(AssertionError): model = DeformConv2d(3, 2, 3, groups=2) # test out_channels % group != 0 with pytest.raises(AssertionError): model = DeformConv2d(3, 4, 3, groups=3)
def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1, dilation=3, groups=1, bias=False, type='dilation', init_cfg=dict(type='Normal', std=0.01, override=dict(name='conv'))): super(AdaptiveConv, self).__init__(init_cfg) assert type in ['offset', 'dilation'] self.adapt_type = type assert kernel_size == 3, 'Adaptive conv only supports kernels 3' if self.adapt_type == 'offset': assert stride == 1 and padding == 1 and groups == 1, \ 'Adaptive conv offset mode only supports padding: {1}, ' \ f'stride: {1}, groups: {1}' self.conv = DeformConv2d(in_channels, out_channels, kernel_size, padding=padding, stride=stride, groups=groups, bias=bias) else: self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, padding=dilation, dilation=dilation)
def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1, dilation=1, deformable_groups=1): super(DeformConvWithOff, self).__init__() self.offset_conv = nn.Conv2d( in_channels, deformable_groups * 2 * kernel_size * kernel_size, kernel_size=kernel_size, stride=stride, padding=padding, ) self.dcn = DeformConv2d( in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, deformable_groups=deformable_groups, )
def _init_layers(self): """Initialize layers of the head.""" self.relu = nn.ReLU(inplace=True) self.cls_convs = nn.ModuleList() self.reg_convs = nn.ModuleList() for i in range(self.stacked_convs): chn = self.in_channels if i == 0 else self.feat_channels self.cls_convs.append( ConvModule( chn, self.feat_channels, 3, stride=1, padding=1, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg)) self.reg_convs.append( ConvModule( chn, self.feat_channels, 3, stride=1, padding=1, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg)) pts_out_dim = 4 if self.use_grid_points else 2 * self.num_points self.reppoints_cls_conv = DeformConv2d(self.feat_channels, self.point_feat_channels, self.dcn_kernel, 1, self.dcn_pad) self.reppoints_cls_out = nn.Conv2d(self.point_feat_channels, self.cls_out_channels, 1, 1, 0) self.reppoints_pts_init_conv = nn.Conv2d(self.feat_channels, self.point_feat_channels, 3, 1, 1) self.reppoints_pts_init_out = nn.Conv2d(self.point_feat_channels, pts_out_dim, 1, 1, 0) self.reppoints_pts_refine_conv = DeformConv2d(self.feat_channels, self.point_feat_channels, self.dcn_kernel, 1, self.dcn_pad) self.reppoints_pts_refine_out = nn.Conv2d(self.point_feat_channels, pts_out_dim, 1, 1, 0)
def __init__(self, in_channels, out_channels, kernel_size=3, deform_groups=4): super(FeatureAdaption, self).__init__() offset_channels = kernel_size * kernel_size * 2 self.conv_offset = nn.Conv2d(34, deform_groups * offset_channels, 1, bias=False) self.conv_adaption_reg = DeformConv2d(in_channels, out_channels, kernel_size=kernel_size, padding=(kernel_size - 1) // 2, deform_groups=deform_groups) self.conv_adaption_cls = DeformConv2d(in_channels, out_channels, kernel_size=kernel_size, padding=(kernel_size - 1) // 2, deform_groups=deform_groups) self.relu = nn.ReLU(inplace=True)
def __init__(self, in_channels, out_channels, kernel_size=3, deform_groups=1): super(AlignConv, self).__init__() self.kernel_size = kernel_size self.deform_conv = DeformConv2d(in_channels, out_channels, kernel_size=kernel_size, padding=(kernel_size - 1) // 2, deform_groups=deform_groups) self.relu = nn.ReLU(inplace=True)
def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1, dilation=1, groups=1, offset_group=1): super(DeformConvWarp, self).__init__() self.DCN_V1 = DeformConv2d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=False)
def __init__(self, in_channels, out_channels, kernel_size=3, deform_groups=4, init_cfg=dict( type='Normal', layer='Conv2d', std=0.1, override=dict( type='Normal', name='conv_adaption', std=0.01))): super(FeatureAdaption, self).__init__(init_cfg) offset_channels = kernel_size * kernel_size * 2 self.conv_offset = nn.Conv2d( 2, deform_groups * offset_channels, 1, bias=False) self.conv_adaption = DeformConv2d( in_channels, out_channels, kernel_size=kernel_size, padding=(kernel_size - 1) // 2, deform_groups=deform_groups) self.relu = nn.ReLU(inplace=True)
def adm_multibox(level_channels, anchor_nums, num_classes): assert set(anchor_nums) == {3} adm_loc_layers1 = [] adm_loc_layers2 = [] adm_loc_layers3 = [] adm_conf_layers1 = [] adm_conf_layers2 = [] adm_conf_layers3 = [] for _ in level_channels: adm_loc_layers1 += [DeformConv2d(256, 4, kernel_size=3, padding=1)] adm_loc_layers2 += [DeformConv2d(256, 4, kernel_size=3, padding=1)] adm_loc_layers3 += [DeformConv2d(256, 4, kernel_size=3, padding=1)] adm_conf_layers1 += [DeformConv2d(256, num_classes, kernel_size=3, padding=1)] adm_conf_layers2 += [DeformConv2d(256, num_classes, kernel_size=3, padding=1)] adm_conf_layers3 += [DeformConv2d(256, num_classes, kernel_size=3, padding=1)] return ( (adm_loc_layers1, adm_loc_layers2, adm_loc_layers3), (adm_conf_layers1, adm_conf_layers2, adm_conf_layers3))
def __init__(self, num_classes=80, in_channels=(512, 1024, 512, 256, 256, 256), anchor_generator=dict( type='SSDAnchorGenerator', scale_major=False, input_size=300, strides=[8, 16, 32, 64, 100, 300], ratios=([2], [2, 3], [2, 3], [2, 3], [2], [2]), basesize_ratio_range=(0.1, 0.9)), background_label=None, bbox_coder=dict( type='DeltaXYWHBBoxCoder', target_means=[.0, .0, .0, .0], target_stds=[1.0, 1.0, 1.0, 1.0], ), reg_decoded_bbox=False, train_cfg=None, test_cfg=None): super(AnchorHead, self).__init__() self.num_classes = num_classes self.in_channels = in_channels self.cls_out_channels = num_classes + 1 # add background class self.anchor_generator = build_anchor_generator(anchor_generator) num_anchors = self.anchor_generator.num_base_anchors self.num_anchors = num_anchors reg_convs = [] cls_convs = [] for i in range(len(in_channels)): reg_convs.append( nn.Conv2d( in_channels[i], num_anchors[i] * 4 , kernel_size=3, padding=1)) cls_convs.append( nn.Conv2d( in_channels[i], num_anchors[i]* (num_classes + 1), kernel_size=3, padding=1)) self.reg_convs = nn.ModuleList(reg_convs) self.cls_convs = nn.ModuleList(cls_convs) self.background_label = ( num_classes if background_label is None else background_label) # background_label should be either 0 or num_classes assert (self.background_label == 0 or self.background_label == num_classes) self.bbox_coder = build_bbox_coder(bbox_coder) self.reg_decoded_bbox = reg_decoded_bbox self.use_sigmoid_cls = False self.cls_focal_loss = False self.train_cfg = train_cfg self.test_cfg = test_cfg # set sampling=False for archor_target self.sampling = False if self.train_cfg: self.assigner = build_assigner(self.train_cfg.assigner) # SSD sampling=False so use PseudoSampler sampler_cfg = dict(type='PseudoSampler') self.sampler = build_sampler(sampler_cfg, context=self) self.fp16_enabled = False # dcn dcn_reg_convs = [] dcn_cls_convs = [] for i in range(len(in_channels)): dcn_reg_convs.append( DeformConv2d( in_channels[i]*num_anchors[i], 4*num_anchors[i], kernel_size=3, bias=True, padding=3, groups= num_anchors[i], deform_groups= num_anchors[i], dilation = 3 )) dcn_cls_convs.append( DeformConv2d( in_channels[i]*num_anchors[i], (num_classes + 1)*num_anchors[i], kernel_size=3, bias=True, padding=3, groups= num_anchors[i], deform_groups= num_anchors[i], dilation=3 )) self.dcn_reg_convs = nn.ModuleList(dcn_reg_convs) self.dcn_cls_convs = nn.ModuleList(dcn_cls_convs) self.BCE = nn.BCEWithLogitsLoss() self.count = 0.0 self.total_count = 3230.0*24.0 self.offset_transfrom = nn.Linear(2,2)
def _test_deformconv(self, dtype=torch.float, threshold=1e-3, device='cuda', batch_size=10, im2col_step=2): if not torch.cuda.is_available() and device == 'cuda': pytest.skip('test requires GPU') from mmcv.ops import DeformConv2dPack c_in = 1 c_out = 1 batch_size = 10 repeated_input = np.repeat(input, batch_size, axis=0) repeated_gt_out = np.repeat(gt_out, batch_size, axis=0) repeated_gt_x_grad = np.repeat(gt_x_grad, batch_size, axis=0) x = torch.tensor(repeated_input, device=device, dtype=dtype) x.requires_grad = True model = DeformConv2dPack(in_channels=c_in, out_channels=c_out, kernel_size=2, stride=1, padding=0, im2col_step=im2col_step) model.conv_offset.weight.data = torch.nn.Parameter( torch.Tensor(offset_weight).reshape(8, 1, 2, 2)) model.conv_offset.bias.data = torch.nn.Parameter( torch.Tensor(offset_bias).reshape(8)) model.weight.data = torch.nn.Parameter( torch.Tensor(deform_weight).reshape(1, 1, 2, 2)) if device == 'cuda': model.cuda() model.type(dtype) out = model(x) out.backward(torch.ones_like(out)) assert np.allclose(out.data.detach().cpu().numpy(), repeated_gt_out, threshold) assert np.allclose(x.grad.detach().cpu().numpy(), repeated_gt_x_grad, threshold) # the batch size of the input is increased which results in # a larger gradient so we need to divide by the batch_size assert np.allclose( model.conv_offset.weight.grad.detach().cpu().numpy() / batch_size, gt_offset_weight_grad, threshold) assert np.allclose( model.conv_offset.bias.grad.detach().cpu().numpy() / batch_size, gt_offset_bias_grad, threshold) assert np.allclose( model.weight.grad.detach().cpu().numpy() / batch_size, gt_deform_weight_grad, threshold) from mmcv.ops import DeformConv2d # test bias model = DeformConv2d(1, 1, 2, stride=1, padding=0) assert not hasattr(model, 'bias') # test bias=True with pytest.raises(AssertionError): model = DeformConv2d(1, 1, 2, stride=1, padding=0, bias=True) # test in_channels % group != 0 with pytest.raises(AssertionError): model = DeformConv2d(3, 2, 3, groups=2) # test out_channels % group != 0 with pytest.raises(AssertionError): model = DeformConv2d(3, 4, 3, groups=3)
def __init__(self, in_channels, out_channels, inner_channels, deform_groups=17, dilations=(3, 6, 12, 18, 24), trans_conv_kernel=1, res_blocks_cfg=None, offsets_kernel=3, deform_conv_kernel=3, in_index=0, input_transform=None, freeze_trans_layer=True, norm_eval=False, im2col_step=80): super().__init__() self.in_channels = in_channels self.out_channels = out_channels self.inner_channels = inner_channels self.deform_groups = deform_groups self.dilations = dilations self.trans_conv_kernel = trans_conv_kernel self.res_blocks_cfg = res_blocks_cfg self.offsets_kernel = offsets_kernel self.deform_conv_kernel = deform_conv_kernel self.in_index = in_index self.input_transform = input_transform self.freeze_trans_layer = freeze_trans_layer self.norm_eval = norm_eval self.im2col_step = im2col_step identity_trans_layer = False assert trans_conv_kernel in [0, 1, 3] kernel_size = trans_conv_kernel if kernel_size == 3: padding = 1 elif kernel_size == 1: padding = 0 else: # 0 for Identity mapping. identity_trans_layer = True if identity_trans_layer: self.trans_layer = nn.Identity() else: self.trans_layer = build_conv_layer(cfg=dict(type='Conv2d'), in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=1, padding=padding) # build chain of residual blocks if res_blocks_cfg is not None and not isinstance(res_blocks_cfg, dict): raise TypeError('res_blocks_cfg should be dict or None.') if res_blocks_cfg is None: block_type = 'BASIC' num_blocks = 20 else: block_type = res_blocks_cfg.get('block', 'BASIC') num_blocks = res_blocks_cfg.get('num_blocks', 20) block = self.blocks_dict[block_type] res_layers = [] downsample = nn.Sequential( build_conv_layer(cfg=dict(type='Conv2d'), in_channels=out_channels, out_channels=inner_channels, kernel_size=1, stride=1, bias=False), build_norm_layer(dict(type='BN'), inner_channels)[1]) res_layers.append( block(in_channels=out_channels, out_channels=inner_channels, downsample=downsample)) for _ in range(1, num_blocks): res_layers.append(block(inner_channels, inner_channels)) self.offset_feats = nn.Sequential(*res_layers) # build offset layers self.num_offset_layers = len(dilations) assert self.num_offset_layers > 0, 'Number of offset layers ' \ 'should be larger than 0.' target_offset_channels = 2 * offsets_kernel**2 * deform_groups offset_layers = [ build_conv_layer( cfg=dict(type='Conv2d'), in_channels=inner_channels, out_channels=target_offset_channels, kernel_size=offsets_kernel, stride=1, dilation=dilations[i], padding=dilations[i], bias=False, ) for i in range(self.num_offset_layers) ] self.offset_layers = nn.ModuleList(offset_layers) # build deformable conv layers assert digit_version(mmcv.__version__) >= \ digit_version(self.minimum_mmcv_version), \ f'Current MMCV version: {mmcv.__version__}, ' \ f'but MMCV >= {self.minimum_mmcv_version} is required, see ' \ f'https://github.com/open-mmlab/mmcv/issues/1440, ' \ f'Please install the latest MMCV.' if has_mmcv_full: deform_conv_layers = [ DeformConv2d( in_channels=out_channels, out_channels=out_channels, kernel_size=deform_conv_kernel, stride=1, padding=int(deform_conv_kernel / 2) * dilations[i], dilation=dilations[i], deform_groups=deform_groups, im2col_step=self.im2col_step, ) for i in range(self.num_offset_layers) ] else: raise ImportError('Please install the full version of mmcv ' 'to use `DeformConv2d`.') self.deform_conv_layers = nn.ModuleList(deform_conv_layers) self.freeze_layers()
def _init_layers(self): """Initialize layers of the head.""" self.relu = nn.ReLU(inplace=True) self.relu2 = nn.LeakyReLU(0.1, inplace=False) self.sigmoid = nn.Sigmoid() self.cls_convs = nn.ModuleList() self.reg_convs = nn.ModuleList() self.cls_domain = nn.ModuleList() self.bce = nn.BCELoss() self.gradreverse = GradReverse() for i in range(self.stacked_convs): chn = self.in_channels if i == 0 else self.feat_channels self.cls_convs.append( ConvModule(chn, self.feat_channels, 3, stride=1, padding=1, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg)) self.reg_convs.append( ConvModule(chn, self.feat_channels, 3, stride=1, padding=1, conv_cfg=self.conv_cfg, norm_cfg=self.norm_cfg)) for i, channels in enumerate( [[self.feat_channels, self.feat_channels], [self.feat_channels, int(self.feat_channels / 2)], [int(self.feat_channels / 2), 1]]): chn_in = channels[0] chn_out = channels[1] self.cls_domain.append( ConvModule(chn_in, chn_out, 1, stride=1, padding=0, conv_cfg=self.conv_cfg, norm_cfg=None, act_cfg=None)) pts_out_dim = 4 if self.use_grid_points else 2 * self.num_points self.reppoints_cls_conv = DeformConv2d(self.feat_channels, self.point_feat_channels, self.dcn_kernel, 1, self.dcn_pad) self.reppoints_cls_out = nn.Conv2d(self.point_feat_channels, self.cls_out_channels, 1, 1, 0) self.reppoints_pts_init_conv = nn.Conv2d(self.feat_channels, self.point_feat_channels, 3, 1, 1) self.reppoints_pts_init_out = nn.Conv2d(self.point_feat_channels, pts_out_dim, 1, 1, 0) self.reppoints_pts_refine_conv = DeformConv2d(self.feat_channels, self.point_feat_channels, self.dcn_kernel, 1, self.dcn_pad) self.reppoints_pts_refine_out = nn.Conv2d(self.point_feat_channels, pts_out_dim, 1, 1, 0)
def test_deform_conv2d(threshold=1e-3): try: from mmcv.ops import DeformConv2d, get_onnxruntime_op_path except (ImportError, ModuleNotFoundError): pytest.skip('deform_conv op is not successfully compiled') ort_custom_op_path = get_onnxruntime_op_path() if not os.path.exists(ort_custom_op_path): pytest.skip('custom ops for onnxruntime are not compiled.') # deform conv config # modulated deform conv config in_channels = 1 out_channels = 64 stride = 1 padding = 0 dilation = 1 groups = 1 deform_groups = 1 kernel_size = 2 input = [[[[1., 2., 3.], [0., 1., 2.], [3., 5., 2.]]]] offset_weight = [[[0.1, 0.4, 0.6, 0.1]], [[0.3, 0.2, 0.1, 0.3]], [[0.5, 0.5, 0.2, 0.8]], [[0.8, 0.3, 0.9, 0.1]], [[0.3, 0.1, 0.2, 0.5]], [[0.3, 0.7, 0.5, 0.3]], [[0.6, 0.2, 0.5, 0.3]], [[0.4, 0.1, 0.8, 0.4]]] offset_bias = [0.7, 0.1, 0.8, 0.5, 0.6, 0.5, 0.4, 0.7] deform_weight = [[[0.4, 0.2, 0.1, 0.9]]] x = torch.tensor(input) conv_offset = nn.Conv2d(in_channels=in_channels, out_channels=deform_groups * 2 * kernel_size * kernel_size, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, bias=True) conv_offset.weight.data = torch.nn.Parameter( torch.Tensor(offset_weight).reshape(8, 1, 2, 2)) conv_offset.bias.data = torch.nn.Parameter( torch.Tensor(offset_bias).reshape(8)) offset = conv_offset(x) model = DeformConv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, deform_groups) model.weight.data = torch.nn.Parameter( torch.Tensor(deform_weight).reshape(1, 1, 2, 2)) with torch.no_grad(): torch.onnx.export(model, (x, offset), onnx_file, export_params=True, keep_initializers_as_inputs=True, input_names=['input', 'offset'], opset_version=11) session_options = rt.SessionOptions() if os.path.exists(ort_custom_op_path): session_options.register_custom_ops_library(ort_custom_op_path) # compute onnx_output sess = rt.InferenceSession(onnx_file, session_options) onnx_output = sess.run( None, { 'input': x.cpu().detach().numpy(), 'offset': offset.cpu().detach().numpy(), })[0] # compute pytorch_output with torch.no_grad(): pytorch_output = model(x, offset).cpu() # allclose assert np.allclose(pytorch_output, onnx_output, atol=1e-3)