def __init__(self, D=4, R=6, C=32): super(Colorizer, self).__init__() self.D = D self.R = R # window size self.C = C self.P = self.R * 2 + 1 self.N = self.P * self.P self.count = 0 self.memory_patch_R = 12 self.memory_patch_P = self.memory_patch_R * 2 + 1 self.memory_patch_N = self.memory_patch_P * self.memory_patch_P self.correlation_sampler_dilated = [ SpatialCorrelationSampler(kernel_size=1, patch_size=self.memory_patch_P, stride=1, padding=0, dilation=1, dilation_patch=dirate) for dirate in range(2, 6) ] self.correlation_sampler = SpatialCorrelationSampler(kernel_size=1, patch_size=self.P, stride=1, padding=0, dilation=1)
def __init__(self): super(Net, self).__init__() self.args = pwcnet_args() args = self.args self.feature_pyramid_extractor = FeaturePyramidExtractor(args).to(args.device) self.warping_layer = WarpingLayer(args) self.corr = SpatialCorrelationSampler(kernel_size=1, patch_size=(args.search_range*2)+1, stride=1).to(args.device) self.flow_estimators = [] for l, ch in enumerate(args.lv_chs[::-1]): layer = OpticalFlowEstimator(args, ch + (args.search_range*2+1)**2 + 2).to(args.device) self.add_module(f'FlowEstimator(Lv{l})', layer) self.flow_estimators.append(layer) if args.context: self.context_networks = [] for l, ch in enumerate(args.lv_chs[::-1]): layer = ContextNetwork(args, ch + 2).to(args.device) self.add_module(f'ContextNetwork(Lv{l})', layer) self.context_networks.append(layer) # init for m in self.modules(): if isinstance(m, nn.Conv2d): if m.bias is not None: nn.init.uniform_(m.bias) nn.init.xavier_uniform_(m.weight) if isinstance(m, nn.ConvTranspose2d): if m.bias is not None: nn.init.uniform_(m.bias) nn.init.xavier_uniform_(m.weight)
def cross_correlate(x, y, max_distance=9): """Efficiently computes the cross correlation of x and y. Optimized implementation using correlation_cost. Note that we do not normalize by the feature dimension. Args: x: Float32 tensor of shape [height, width, feature_dim]. y: Float32 tensor of shape [height, width, feature_dim]. max_distance: Integer, the maximum distance in pixel coordinates per dimension which is considered to be in the search window. Returns: Float32 tensor of shape [height, width, (2 * max_distance + 1) ** 2]. """ #corr_op=orrelation(pad_size=max_distance, kernel_size=1, max_displacement=max_distance, stride1=1, stride2=1, corr_multiply=1) corr_op = SpatialCorrelationSampler(kernel_size=1, patch_size=2 * max_distance + 1, stride=1, dilation_patch=1, padding=0) xs = x.permute(2, 0, 1) xs = torch.unsqueeze(xs, 0) ys = y.permute(2, 0, 1) ys = torch.unsqueeze(ys, 0) corr = corr_op(xs, ys) bs, _, _, hh, ww = corr.size() corr = corr.view(bs, -1, hh, ww) corr = torch.squeeze(corr, 0) corr = corr.permute(1, 2, 0) # feature_dim=x.size()[-1] # corr *= feature_dim return corr
def correlation1d_cost( reference_fm, target_fm, max_disp=192, start_disp=0, dilation=1, disp_sample=None, kernel_size=1, stride=1, padding=0, dilation_patch=1, ): # for a pixel of left image at (x, y), it will calculates correlation cost volume # with pixel of right image at (xr, y), where xr in [x-max_disp, x+max_disp] # but we only need the left half part, i.e., [x-max_disp, 0] correlation_sampler = SpatialCorrelationSampler( patch_size=(1, max_disp * 2 - 1), kernel_size=kernel_size, stride=stride, padding=padding, dilation_patch=dilation_patch) # [B, 1, max_disp*2-1, H, W] out = correlation_sampler(reference_fm, target_fm) # [B, max_disp*2-1, H, W] out = out.squeeze(1) # [B, max_disp, H, W], grad the left half searching part out = out[:, :max_disp, :, :] cost = F.leaky_relu(out, negative_slope=0.1, inplace=True) return cost
def __init__(self): super(FlowNet, self).__init__() print('FlowNet...') self.debug = False # self.debug = True self.heatmap_size = hyp.flow_patch_size # self.compress_dim = 8 # self.scales = [0.25, 0.5, 1.0] self.scales = [0.125, 0.25, 0.5, 0.75, 1.0] self.num_scales = len(self.scales) # # slightly diff from flownet, here i am using one set of params across all scales # self.compressor = nn.Sequential( # nn.Conv3d(in_channels=32, out_channels=self.compress_dim, kernel_size=1, stride=1, padding=0), # ) self.correlation_sampler = SpatialCorrelationSampler( kernel_size=1, patch_size=self.heatmap_size, stride=1, padding=0, dilation_patch=1, ) self.flow_predictor = nn.Sequential( nn.Conv3d(in_channels=(self.heatmap_size**3), out_channels=64, kernel_size=1, stride=1, padding=0), nn.LeakyReLU(), nn.Conv3d(in_channels=64, out_channels=3, kernel_size=1, stride=1, padding=0), ) self.smoothl1 = torch.nn.SmoothL1Loss(reduction='none') self.smoothl1_mean = torch.nn.SmoothL1Loss(reduction='mean')
def corr(self, refimg_fea, targetimg_fea, maxdisp, fac=1): """ correlation function. Adopted from https://github.com/ClementPinard/Pytorch-Correlation-extension faster, but backwards not implemented. """ from spatial_correlation_sampler import SpatialCorrelationSampler corr = SpatialCorrelationSampler(kernel_size=1,patch_size=(int(1+2*maxdisp//fac),int(1+2*maxdisp)),stride=1,padding=0,dilation_patch=1) cost = corr(refimg_fea, targetimg_fea) cost = F.leaky_relu(cost, 0.1,inplace=True) return cost
def __init__(self): super(FlowNet, self).__init__() print('FlowNet...') self.debug = False # self.debug = True self.heatmap_size = hyp.flow_heatmap_size self.scales = [1.0] self.num_scales = len(self.scales) assert (self.num_scales == 1 ) # do not touch scales, unless you adjust the masking dilation = 1 grid_z, grid_y, grid_x = utils_basic.meshgrid3D( 1, self.heatmap_size, self.heatmap_size, self.heatmap_size) self.max_disp = int(dilation * (self.heatmap_size - 1) / 2) if self.debug: print('max_disp', self.max_disp) self.grid = torch.stack([grid_z, grid_y, grid_x], dim=1) - int( self.heatmap_size / 2) # this is 1 x 3 x H x H x H, with 0 in the middle self.grid = self.grid.reshape(1, 3, -1, 1, 1, 1) * dilation # now we are ready to mult with the cc out # self.compress_dim = 16 # self.compressor = nn.Sequential( # nn.Conv3d(in_channels=hyp.feat_dim, out_channels=self.compress_dim, kernel_size=1, stride=1, padding=0), # ) self.correlation_sampler = SpatialCorrelationSampler( kernel_size=1, patch_size=self.heatmap_size, stride=1, padding=0, dilation_patch=dilation, ).cuda() # self.flow_predictor = nn.Sequential( # nn.Conv3d(in_channels=(self.heatmap_size**3), out_channels=64, kernel_size=3, stride=1, padding=1), # nn.LeakyReLU(negative_slope=0.1), # nn.Conv3d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1), # nn.LeakyReLU(negative_slope=0.1), # nn.Conv3d(in_channels=64, out_channels=3, kernel_size=1, stride=1, padding=0), # ).cuda() self.smoothl1 = torch.nn.SmoothL1Loss(reduction='none') self.smoothl1_mean = torch.nn.SmoothL1Loss(reduction='mean') self.mse = torch.nn.MSELoss(reduction='none') self.mse_mean = torch.nn.MSELoss(reduction='mean')
def __init__(self): super(Colorizer, self).__init__() self.D = 4 self.R = 6 # window size self.C = 16 self.P = self.R * 2 + 1 self.correlation_sampler = SpatialCorrelationSampler(kernel_size=1, patch_size=self.P, stride=1, padding=0, dilation=1)
def __init__(self, D=4, R=6, C=32, mode='faster', training=False, ksargmax=True): super(Colorizer, self).__init__() self.D = D self.R = R # window size self.C = C self.P = self.R * 2 + 1 self.N = self.P * self.P self.count = 0 self.training = training self.mode = mode self.beta = 50 self.kernel_sigma = 1.0 self.ksargmax = ksargmax self.memory_patch_R = 12 self.memory_patch_P = self.memory_patch_R * 2 + 1 self.memory_patch_N = self.memory_patch_P * self.memory_patch_P self.correlation_sampler_dilated = [ SpatialCorrelationSampler(kernel_size=1, patch_size=self.memory_patch_P, stride=1, padding=0, dilation=1, dilation_patch=dirate) for dirate in range(2, 6) ] self.correlation_sampler = SpatialCorrelationSampler(kernel_size=1, patch_size=self.P, stride=1, padding=0, dilation=1)
def __init__(self): super(CorrelationHead, self).__init__() representation_size = 1024 feature_map_res = 7 correlation_patch_size = 16 # flownet uses 21, if we want to use faster-rcnn weight needs to be 16 self.name = "CorrelationHead" self.roi_heads = None self.correlation_layer = SpatialCorrelationSampler( patch_size=correlation_patch_size, dilation_patch=2) self.fc1 = nn.Linear(correlation_patch_size**2 * feature_map_res**2, representation_size) self.fc2 = nn.Linear(representation_size, representation_size) self.fc3 = nn.Linear(representation_size, 4)
def __init__(self, ks, patch, stride, pad, patch_dilation): super(Matching_layer, self).__init__() self.relu = nn.ReLU() self.patch = patch self.correlation_sampler = SpatialCorrelationSampler(ks, patch, stride, pad, patch_dilation)
def local_previous_frame_nearest_neighbor_features_per_object( prev_frame_embedding, query_embedding, prev_frame_labels, gt_ids, max_distance=15): """Computes nearest neighbor features while only allowing local matches. Args: prev_frame_embedding: Tensor of shape [height, width, embedding_dim], the embedding vectors for the last frame. query_embedding: Tensor of shape [height, width, embedding_dim], the embedding vectors for the query frames. prev_frame_labels: Tensor of shape [height, width, 1], the class labels of the previous frame. gt_ids: Int Tensor of shape [n_objs] of the sorted unique ground truth ids in the first frame. max_distance: Integer, the maximum distance allowed for local matching. Returns: nn_features: A float32 np.array of nearest neighbor features of shape [1, height, width, n_objects, 1]. """ if USE_CORRELATION_COST: d = local_pairwise_distances(query_embedding, prev_frame_embedding, max_distance=max_distance) else: d = local_pairwise_distances2(query_embedding, prev_frame_embedding, max_distance=max_distance) # d = (torch.sigmoid(d) - 0.5) * 2 height, width = prev_frame_embedding.size()[:2] if USE_CORRELATION_COST: # corr_op=Correlation(pad_size=max_distance, kernel_size=1, max_displacement=max_distance, stride1=1, stride2=1, corr_multiply=1) corr_op = SpatialCorrelationSampler(kernel_size=1, patch_size=2 * max_distance + 1, stride=1, dilation_patch=1, padding=0) # New, faster code with cross-correlation via correlation_cost. # Due to padding we have to add 1 to the labels. tmp_prev_frame_labels = (prev_frame_labels + 1).float().permute( 2, 0, 1) tmp_prev_frame_labels = torch.unsqueeze(tmp_prev_frame_labels, 0) ones_ = torch.ones_like(tmp_prev_frame_labels) offset_labels = corr_op(ones_, tmp_prev_frame_labels) bs, _, _, hh, ww = offset_labels.size() offset_labels = offset_labels.view(bs, -1, hh, ww) offset_labels = torch.squeeze(offset_labels, 0) offset_labels = offset_labels.permute(1, 2, 0) offset_labels = torch.unsqueeze(offset_labels, 3) offset_labels = torch.round(offset_labels - 1) offset_masks = torch.eq( offset_labels, gt_ids.float().unsqueeze(0).unsqueeze(0).unsqueeze(0)) else: masks = torch.eq(prev_frame_labels, gt_ids.unsqueeze(0).unsqueeze(0)) padded_masks = nn.functional.pad(masks, ( 0, 0, max_distance, max_distance, max_distance, max_distance, )) offset_masks = [] for y_start in range(2 * max_distance + 1): y_end = y_start + height masks_slice = padded_masks[y_start:y_end] for x_start in range(2 * max_distance + 1): x_end = x_start + width offset_mask = masks_slice[:, x_start:x_end] offset_masks.append(offset_mask) offset_masks = torch.stack(offset_masks, dim=2) # pad = torch.ones((height, width, (2 * max_distance + 1) ** 2, gt_ids.size(0))) d_tiled = d.unsqueeze(-1).repeat((1, 1, 1, gt_ids.size(0))) pad = torch.ones_like(d_tiled) # if torch.cuda.is_available(): # pad=pad.cuda() # d_tiled = d_tiled.cuda() d_masked = torch.where(offset_masks, d_tiled, pad) dists, _ = torch.min(d_masked, dim=2) dists = dists.view(1, height, width, gt_ids.size(0), 1) return dists
parser.add_argument('-k', '--kernel-size', type=int, default=3) parser.add_argument('--patch', type=int, default=3) parser.add_argument('--patch_dilation', type=int, default=2) parser.add_argument('-c', '--channel', type=int, default=10) parser.add_argument('--height', type=int, default=10) parser.add_argument('-w', '--width', type=int, default=10) parser.add_argument('-s', '--stride', type=int, default=2) parser.add_argument('-p', '--pad', type=int, default=5) parser.add_argument('-v', '--verbose', action='store_true') args = parser.parse_args() assert (torch.cuda.is_available()), "no comparison to make" device = torch.device("cuda") input1 = torch.randn(args.batch_size, args.channel, args.height, args.width).double() input2 = torch.randn(args.batch_size, args.channel, args.height, args.width).double() input1.requires_grad = True input2.requires_grad = True correlation_sampler = SpatialCorrelationSampler(args.kernel_size, args.patch, args.stride, args.pad, args.patch_dilation) if 'forward' in args.direction: check_forward(input1, input2, correlation_sampler, args.verbose) if 'backward' in args.direction: check_backward(input1, input2, correlation_sampler, args.verbose)
def __init__(self, use_norm=True, num_class=2, layer_nums=(3, 5, 5), layer_strides=(2, 2, 2), num_filters=(128, 128, 256), upsample_strides=(1, 2, 4), num_upsample_filters=(256, 256, 256), num_input_features=128, num_anchor_per_loc=2, encode_background_as_zeros=True, use_direction_classifier=True, use_groupnorm=False, num_groups=32, box_code_size=7, num_direction_bins=2, corr_patch_size=9, corr_kernel_size=3, corr_dilation_patch=1, voting_range=1, name='rpn'): """upsample_strides support float: [0.25, 0.5, 1] if upsample_strides < 1, conv2d will be used instead of convtranspose2d. """ super(RPNBase_tracking, self).__init__( use_norm=use_norm, num_class=num_class, layer_nums=layer_nums, layer_strides=layer_strides, num_filters=num_filters, upsample_strides=upsample_strides, num_upsample_filters=num_upsample_filters, num_input_features=num_input_features, num_anchor_per_loc=num_anchor_per_loc, encode_background_as_zeros=encode_background_as_zeros, use_direction_classifier=use_direction_classifier, use_groupnorm=use_groupnorm, num_groups=num_groups, box_code_size=box_code_size, num_direction_bins=num_direction_bins, corr_patch_size=corr_patch_size, corr_kernel_size=corr_kernel_size, corr_dilation_patch=corr_dilation_patch, voting_range=voting_range, name=name) self._num_anchor_per_loc = num_anchor_per_loc self._num_direction_bins = num_direction_bins self._num_class = num_class self._use_direction_classifier = use_direction_classifier self._box_code_size = box_code_size self._corr_patch_size = corr_patch_size self._corr_kernel_size = corr_kernel_size self._corr_dilation_patch = corr_dilation_patch self._voting_range = voting_range if encode_background_as_zeros: num_cls = num_anchor_per_loc * num_class else: num_cls = num_anchor_per_loc * (num_class + 1) if len(num_upsample_filters) == 0: final_num_filters = self._num_out_filters else: final_num_filters = sum(num_upsample_filters) self.conv_cls = nn.Conv2d(final_num_filters * 2, num_cls, 1) self.conv_box = nn.Conv2d(final_num_filters * 2, num_anchor_per_loc * box_code_size, 1) if use_direction_classifier: self.conv_dir_cls = nn.Conv2d( final_num_filters * 2, num_anchor_per_loc * num_direction_bins, 1) ''' - Output sizes oH and oW are no longer dependant of patch size, but only of kernel size and padding - Patch size patch_size is now the whole patch, and not only the radii. - stride1 is now stride - stride2 is dilation_patch, which behave like dilated convolutions - equivalent max_displacement is then dilation_patch * (patch_size - 1) / 2. ''' self.correlation_sampler = SpatialCorrelationSampler( kernel_size=self._corr_kernel_size, patch_size=self._corr_patch_size, stride=1, # padding=1 + self._corr_kernel_size // 3, padding=2, dilation_patch=self._corr_dilation_patch) self.resample = Resample2d()
def __init__(self, md=4): """ input: md --- maximum displacement (for correlation. default: 4), after warpping """ super(FlowNetPwcLite2xA, self).__init__() # self.conv1a = conv(1, 8, kernel_size=3, stride=2) # self.conv1aa = conv(8, 8, kernel_size=3, stride=1) # self.conv1b = conv(8, 8, kernel_size=3, stride=1) self.conv2a = conv(1, 8, kernel_size=3, stride=2) self.conv2aa = conv(8, 8, kernel_size=3, stride=1) self.conv2b = conv(8, 8, kernel_size=3, stride=1) self.conv3a = conv(8, 32, kernel_size=3, stride=2) self.conv3aa = conv(32, 32, kernel_size=3, stride=1) self.conv3b = conv(32, 32, kernel_size=3, stride=1) self.conv4a = conv(32, 64, kernel_size=3, stride=2) self.conv4aa = conv(64, 64, kernel_size=3, stride=1) self.conv4b = conv(64, 32, kernel_size=3, stride=1) self.conv5a = conv(32, 64, kernel_size=3, stride=2) self.conv5aa = conv(64, 64, kernel_size=3, stride=1) self.conv5b = conv(64, 32, kernel_size=3, stride=1) self.conv6aa = conv(32, 64, kernel_size=3, stride=2) self.conv6a = conv(64, 64, kernel_size=3, stride=1) self.conv6b = conv(64, 64, kernel_size=3, stride=1) #self.corr = Correlation(pad_size=md, kernel_size=1, max_displacement=md, stride1=1, stride2=1, corr_multiply=1) md2plus1 = md * 2 + 1 self.corr6 = SpatialCorrelationSampler(patch_size=md2plus1, kernel_size=1, stride=1) md2plus1N = (md - 1) * 2 + 1 self.corrN = SpatialCorrelationSampler(patch_size=md2plus1N, kernel_size=1, stride=1) self.leakyRELU = nn.LeakyReLU(0.1) nd = md2plus1**2 dd = np.cumsum([128, 128, 96, 64, 32]) od = nd self.conv6_0 = conv(od, 128, kernel_size=3, stride=1) self.conv6_1 = conv(od + dd[0], 128, kernel_size=3, stride=1) self.conv6_2 = conv(od + dd[1], 96, kernel_size=3, stride=1) self.conv6_3 = conv(od + dd[2], 64, kernel_size=3, stride=1) self.conv6_4 = conv(od + dd[3], 32, kernel_size=3, stride=1) self.predict_flow6 = predict_flow(od + dd[4]) self.deconv6 = deconv(2, 2, kernel_size=4, stride=2, padding=1) self.upfeat6 = deconv(od + dd[4], 16, kernel_size=4, stride=2, padding=1) nd = md2plus1N**2 od = nd + 32 + 2 + 16 self.conv5_0 = conv(od, 128, kernel_size=3, stride=1) self.conv5_1 = conv(od + dd[0], 128, kernel_size=3, stride=1) self.conv5_2 = conv(od + dd[1], 96, kernel_size=3, stride=1) self.conv5_3 = conv(od + dd[2], 64, kernel_size=3, stride=1) self.conv5_4 = conv(od + dd[3], 32, kernel_size=3, stride=1) self.predict_flow5 = predict_flow(od + dd[4]) self.deconv5 = deconv(2, 2, kernel_size=4, stride=2, padding=1) self.upfeat5 = deconv(od + dd[4], 16, kernel_size=4, stride=2, padding=1) od = nd + 32 + 2 + 16 self.conv4_0 = conv(od, 128, kernel_size=3, stride=1) self.conv4_1 = conv(od + dd[0], 128, kernel_size=3, stride=1) self.conv4_2 = conv(od + dd[1], 96, kernel_size=3, stride=1) self.conv4_3 = conv(od + dd[2], 64, kernel_size=3, stride=1) self.conv4_4 = conv(od + dd[3], 32, kernel_size=3, stride=1) self.predict_flow4 = predict_flow(od + dd[4]) self.deconv4 = deconv(2, 2, kernel_size=4, stride=2, padding=1) self.upfeat4 = deconv(od + dd[4], 16, kernel_size=4, stride=2, padding=1) od = nd + 32 + 2 + 16 self.conv3_0 = conv(od, 128, kernel_size=3, stride=1) self.conv3_1 = conv(od + dd[0], 128, kernel_size=3, stride=1) self.conv3_2 = conv(od + dd[1], 96, kernel_size=3, stride=1) self.conv3_3 = conv(od + dd[2], 64, kernel_size=3, stride=1) self.conv3_4 = conv(od + dd[3], 32, kernel_size=3, stride=1) self.predict_flow3 = predict_flow(od + dd[4]) self.deconv3 = deconv(2, 2, kernel_size=4, stride=2, padding=1) self.upfeat3 = deconv(od + dd[4], 16, kernel_size=4, stride=2, padding=1) od = nd + 8 + 2 + 16 self.conv2_0 = conv(od, 128, kernel_size=3, stride=1) self.conv2_1 = conv(od + dd[0], 128, kernel_size=3, stride=1) self.conv2_2 = conv(od + dd[1], 96, kernel_size=3, stride=1) self.conv2_3 = conv(od + dd[2], 64, kernel_size=3, stride=1) self.conv2_4 = conv(od + dd[3], 32, kernel_size=3, stride=1) self.predict_flow2 = predict_flow(od + dd[4]) self.deconv2 = deconv(2, 2, kernel_size=4, stride=2, padding=1) self.dc_conv1 = conv(od + dd[4], 128, kernel_size=3, stride=1, padding=1, dilation=1) self.dc_conv2 = conv(128, 128, kernel_size=3, stride=1, padding=2, dilation=2) self.dc_conv3 = conv(128, 128, kernel_size=3, stride=1, padding=4, dilation=4) self.dc_conv4 = conv(128, 96, kernel_size=3, stride=1, padding=8, dilation=8) self.dc_conv5 = conv(96, 64, kernel_size=3, stride=1, padding=16, dilation=16) self.dc_conv6 = conv(64, 32, kernel_size=3, stride=1, padding=1, dilation=1) self.dc_conv7 = predict_flow(32) for m in self.modules(): if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d): nn.init.kaiming_normal_(m.weight.data, mode='fan_in') if m.bias is not None: m.bias.data.zero_()
elif args.dtype == 'float': dtype = torch.float32 else: dtype = torch.float64 input1 = torch.randn(args.batch_size, args.channel, args.height, args.width, dtype=dtype, device=device, requires_grad=True) input2 = torch.randn_like(input1) correlation_sampler = SpatialCorrelationSampler(args.kernel_size, args.patch, args.stride, args.pad, args.dilation, args.patch_dilation) # Force CUDA initialization output = correlation_sampler(input1, input2) print(output.size()) output.mean().backward() forward_min = float('inf') forward_time = 0 backward_min = float('inf') backward_time = 0 for _ in trange(args.runs): correlation_sampler.zero_grad() start = time.time() output = correlation_sampler(input1, input2)
def __init__(self, use_loc=False, use_trk=False): super(STNNet, self).__init__() self.loc = use_loc self.trk = use_trk # backbone self.backbone = BaseNet() self.relu = nn.ReLU() # multi-scale combination self.deconv1 = nn.ConvTranspose2d(512, 256, 3, stride=2, padding=1, output_padding=1, dilation=2) self.deconv2 = nn.ConvTranspose2d(256, 128, 3, stride=2, padding=1, output_padding=1, dilation=2) self.conv1 = nn.Conv2d(512, 256, 1) self.conv2 = nn.Conv2d(256, 256, 3, padding=1, dilation=1) self.conv3 = nn.Conv2d(256, 128, 1) self.conv4 = nn.Conv2d(128, 128, 3, padding=1, dilation=1) # attention layers self.s_weight1 = SpatialWeightLayer() self.s_weight2 = SpatialWeightLayer() self.s_weight3 = SpatialWeightLayer() # density output layers self.den_output_layer1 = nn.Conv2d(128, 1, kernel_size=1) self.den_output_layer2 = nn.Conv2d(256, 1, kernel_size=1) self.den_output_layer3 = nn.Conv2d(512, 1, kernel_size=1) # fix above parameters when training the association heads if self.trk: for para in self.parameters(): para.requires_grad = False if self.loc: # localization output layers self.loc_output_layer1 = nn.Conv2d(128, 2, 3, padding=1, dilation=1) self.loc_output_layer2 = nn.Conv2d(256, 2, 3, padding=1, dilation=1) self.loc_output_layer3 = nn.Conv2d(512, 2, 3, padding=1, dilation=1) self.reg_output_layer1 = nn.Conv2d(128, 2, 3, padding=1, dilation=1) self.reg_output_layer2 = nn.Conv2d(256, 2, 3, padding=1, dilation=1) self.reg_output_layer3 = nn.Conv2d(512, 2, 3, padding=1, dilation=1) self.s_weight_loc = nn.Conv2d(6, 6, 3, padding=1) self.merge_loc = nn.Conv2d(6, 2, 1, padding=0) self.s_weight_reg = nn.Conv2d(6, 6, 3, padding=1) self.merge_reg = nn.Conv2d(6, 2, 1, padding=0) if self.trk: # tracking output layers self.corr_layer1 = SpatialCorrelationSampler(1, 11, 1, 0, 1) self.corr_layer2 = SpatialCorrelationSampler(1, 11, 1, 0, 1) self.corr_layer3 = SpatialCorrelationSampler(1, 11, 1, 0, 1) self.trk_output_layer1 = nn.Conv2d(363, 128, 3, padding=1, dilation=1) self.trk_output_layer2 = nn.Conv2d(128, 64, 3, padding=1, dilation=1) # graph layers self.gcn_layers = PointConvDensityClsSsg(num_classes=2, num_pt=128) # load weights of the backbone network mod = models.vgg16(pretrained=True) self._initialize_weights() my_models = self.backbone.state_dict() pre_models = list(mod.state_dict().items()) count = 0 for layer_name, value in my_models.items(): prelayer_name, pre_weights = pre_models[count] my_models[layer_name] = pre_weights count += 1 self.backbone.load_state_dict(my_models)