├── src ├── __init__.py ├── AWP │ ├── __init__.py │ └── awp.py ├── utils │ ├── __init__.py │ ├── count_model_params.py │ ├── timer.py │ ├── gpu_memory.py │ ├── model_sl.py │ ├── pad_tensor.py │ ├── dup_stdout_manager.py │ ├── print_easydict.py │ ├── utils_awp.py │ ├── data_to_cuda.py │ └── parse_args.py ├── qap_solvers │ ├── __init__.py │ ├── spectral_matching.py │ ├── rrwm.py │ └── rrwhm.py ├── lap_solvers │ ├── __init__.py │ └── hungarian.py ├── parallel │ ├── __init__.py │ ├── data_parallel.py │ └── scatter_gather.py ├── sparse_torch │ └── __init__.py ├── dataset │ ├── base_dataset.py │ ├── __init__.py │ ├── dataset_config.py │ ├── qaplib.py │ ├── imc_pt_sparsegm.py │ └── willow_obj.py ├── rules.py ├── extension │ ├── sparse_dot │ │ ├── csr_dot_diag_cuda.cu │ │ ├── csr_dot_csc_cuda.cu │ │ └── sparse_dot.cpp │ └── bilinear_diag │ │ └── bilinear_diag_cuda.cu ├── displacement_layer.py ├── backbone.py ├── factorize_graph_matching.py ├── gconv.py ├── build_graphs.py └── spectral_clustering.py ├── .gitignore ├── figures └── attack.png ├── data └── PascalVOC │ └── voc2011_pairs.npz ├── models ├── CIE │ ├── model_config.py │ ├── README.md │ └── model.py ├── GMN │ ├── model_config.py │ ├── voting_layer.py │ ├── displacement_layer.py │ ├── affinity_layer.py │ ├── README.md │ └── model.py ├── PCA │ ├── model_config.py │ ├── README.md │ ├── model.py │ └── affinity_layer.py ├── BBGM │ ├── model_config.py │ ├── affinity_layer.py │ ├── sconv_archs.py │ ├── README.md │ └── model.py ├── GANN │ ├── model_config.py │ └── README.md └── NGM │ ├── model_config.py │ ├── geo_edge_feature.py │ ├── hypermodel.py │ └── model.py ├── experiments ├── eval.yaml ├── eval_blackbox.yaml ├── config1.yaml ├── config2.yaml └── config3.yaml └── README.md /src/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /src/AWP/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /src/utils/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | __pycache__/ 2 | -------------------------------------------------------------------------------- /src/qap_solvers/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /src/lap_solvers/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | Solvers for linear assignment problem (LAP) 3 | """ 4 | -------------------------------------------------------------------------------- /src/parallel/__init__.py: -------------------------------------------------------------------------------- 1 | from .data_parallel import * 2 | 3 | __all__ = ['DataParallel'] 4 | -------------------------------------------------------------------------------- /figures/attack.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Thinklab-SJTU/robustMatch/HEAD/figures/attack.png -------------------------------------------------------------------------------- /data/PascalVOC/voc2011_pairs.npz: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Thinklab-SJTU/robustMatch/HEAD/data/PascalVOC/voc2011_pairs.npz -------------------------------------------------------------------------------- /src/sparse_torch/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | sparse matrix in pytorch implementation. 3 | """ 4 | 5 | from .csx_matrix import CSRMatrix3d, CSCMatrix3d, dot, concatenate 6 | -------------------------------------------------------------------------------- /src/dataset/base_dataset.py: -------------------------------------------------------------------------------- 1 | class BaseDataset: 2 | def __init__(self): 3 | pass 4 | 5 | def get_pair(self, cls, shuffle): 6 | raise NotImplementedError 7 | -------------------------------------------------------------------------------- /src/utils/count_model_params.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | def count_parameters(model): 4 | return np.sum(np.prod(v.size()) for name, v in model.named_parameters() if "auxiliary" not in name) 5 | -------------------------------------------------------------------------------- /src/dataset/__init__.py: -------------------------------------------------------------------------------- 1 | from .pascal_voc import PascalVOC 2 | from .willow_obj import WillowObject 3 | from .qaplib import QAPLIB 4 | from .cub2011 import CUB2011 5 | from .imc_pt_sparsegm import IMC_PT_SparseGM 6 | from .dataset_config import dataset_cfg 7 | -------------------------------------------------------------------------------- /models/CIE/model_config.py: -------------------------------------------------------------------------------- 1 | from easydict import EasyDict as edict 2 | 3 | __C = edict() 4 | 5 | model_cfg = __C 6 | 7 | # CIE model options 8 | __C.CIE = edict() 9 | __C.CIE.FEATURE_CHANNEL = 512 10 | __C.CIE.SK_ITER_NUM = 20 11 | __C.CIE.SK_EPSILON = 1.0e-10 12 | __C.CIE.SK_TAU = 0.005 13 | __C.CIE.GNN_LAYER = 5 14 | __C.CIE.GNN_FEAT = 1024 15 | __C.CIE.CROSS_ITER = False 16 | -------------------------------------------------------------------------------- /models/GMN/model_config.py: -------------------------------------------------------------------------------- 1 | from easydict import EasyDict as edict 2 | 3 | __C = edict() 4 | 5 | model_cfg = __C 6 | 7 | # GMN model options 8 | __C.GMN = edict() 9 | __C.GMN.FEATURE_CHANNEL = 512 10 | __C.GMN.PI_ITER_NUM = 50 11 | __C.GMN.PI_STOP_THRESH = 2e-7 12 | __C.GMN.BS_ITER_NUM = 10 13 | __C.GMN.BS_EPSILON = 1e-10 14 | __C.GMN.VOTING_ALPHA = 2e8 15 | __C.GMN.GM_SOLVER = 'SM' -------------------------------------------------------------------------------- /models/PCA/model_config.py: -------------------------------------------------------------------------------- 1 | from easydict import EasyDict as edict 2 | 3 | __C = edict() 4 | 5 | model_cfg = __C 6 | 7 | # PCA model options 8 | __C.PCA = edict() 9 | __C.PCA.FEATURE_CHANNEL = 512 10 | __C.PCA.SK_ITER_NUM = 20 11 | __C.PCA.SK_EPSILON = 1.0e-10 12 | __C.PCA.SK_TAU = 0.005 13 | __C.PCA.GNN_LAYER = 5 14 | __C.PCA.GNN_FEAT = 1024 15 | __C.PCA.CROSS_ITER = False 16 | __C.PCA.CROSS_ITER_NUM = 1 17 | -------------------------------------------------------------------------------- /models/BBGM/model_config.py: -------------------------------------------------------------------------------- 1 | from easydict import EasyDict as edict 2 | 3 | __C = edict() 4 | 5 | model_cfg = __C 6 | 7 | # BBGM model options 8 | __C.BBGM = edict() 9 | __C.BBGM.SOLVER_NAME = 'LPMP' 10 | __C.BBGM.LAMBDA_VAL = 80.0 11 | __C.BBGM.SOLVER_PARAMS = edict() 12 | __C.BBGM.SOLVER_PARAMS.timeout = 1000 13 | __C.BBGM.SOLVER_PARAMS.primalComputationInterval = 10 14 | __C.BBGM.SOLVER_PARAMS.maxIter = 100 15 | __C.BBGM.FEATURE_CHANNEL = 1024 -------------------------------------------------------------------------------- /src/utils/timer.py: -------------------------------------------------------------------------------- 1 | from time import time 2 | 3 | 4 | class Timer: 5 | def __init__(self): 6 | self.t = time() 7 | self.tk = False 8 | 9 | def tick(self): 10 | self.t = time() 11 | self.tk = True 12 | 13 | def toc(self, tick_again=False): 14 | if not self.tk: 15 | raise RuntimeError('not ticked yet!') 16 | self.tk = False 17 | before_t = self.t 18 | cur_t = time() 19 | if tick_again: 20 | self.t = cur_t 21 | self.tk = True 22 | return cur_t - before_t 23 | -------------------------------------------------------------------------------- /src/parallel/data_parallel.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | from .scatter_gather import scatter_kwargs, gather 3 | 4 | 5 | class DataParallel(nn.DataParallel): 6 | """ 7 | DataParallel wrapper with customized scatter/gather functions 8 | """ 9 | def __init__(self, *args, **kwargs): 10 | super(DataParallel, self).__init__(*args, **kwargs) 11 | 12 | def scatter(self, inputs, kwargs, device_ids): 13 | return scatter_kwargs(inputs, kwargs, device_ids, dim=self.dim) 14 | 15 | def gather(self, outputs, output_device): 16 | return gather(outputs, output_device, dim=self.dim) 17 | -------------------------------------------------------------------------------- /src/utils/gpu_memory.py: -------------------------------------------------------------------------------- 1 | from pynvml import * 2 | nvmlInit() 3 | import torch 4 | 5 | def gpu_free_memory(device_id): 6 | """ 7 | Return total amount of available memory in Bytes 8 | :param device_id: GPU device id (int) 9 | :return: total amount of available memory in Bytes 10 | """ 11 | #ree = nvmlDeviceGetMemoryInfo(nvmlDeviceGetHandleByIndex(device_id)).free 12 | #rsvd = torch.cuda.memory_reserved(device_id) 13 | #used = torch.cuda.memory_allocated(device_id) 14 | torch.cuda.empty_cache() 15 | free = nvmlDeviceGetMemoryInfo(nvmlDeviceGetHandleByIndex(device_id)).free 16 | return free -------------------------------------------------------------------------------- /models/GANN/model_config.py: -------------------------------------------------------------------------------- 1 | from easydict import EasyDict as edict 2 | 3 | __C = edict() 4 | 5 | model_cfg = __C 6 | 7 | # GANN model options 8 | __C.GANN = edict() 9 | __C.GANN.FEATURE_CHANNEL = 1024 10 | __C.GANN.SK_ITER_NUM = 20 11 | __C.GANN.SK_TAU = 0.05 12 | __C.GANN.SK_EPSILON = 1e-10 13 | __C.GANN.UNIV_SIZE = 10 14 | __C.GANN.CLUSTER_ITER = 10 15 | __C.GANN.MGM_ITER = [200, 200] 16 | __C.GANN.INIT_TAU = [0.5, 0.5] 17 | __C.GANN.GAMMA = 0.5 18 | __C.GANN.BETA = [1., 0.] 19 | __C.GANN.CONVERGE_TOL = 1e-5 20 | __C.GANN.MIN_TAU = [1e-2, 1e-2] 21 | __C.GANN.SCALE_FACTOR = 1. 22 | __C.GANN.QUAD_WEIGHT = 1. 23 | __C.GANN.CLUSTER_QUAD_WEIGHT = 1. 24 | __C.GANN.PROJECTOR = ['sinkhorn', 'sinkhorn'] 25 | __C.GANN.NORM_QUAD_TERM = False 26 | -------------------------------------------------------------------------------- /models/BBGM/affinity_layer.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | 5 | class InnerProductWithWeightsAffinity(nn.Module): 6 | def __init__(self, input_dim, output_dim): 7 | super(InnerProductWithWeightsAffinity, self).__init__() 8 | self.d = output_dim 9 | self.A = torch.nn.Linear(input_dim, output_dim) 10 | 11 | def _forward(self, X, Y, weights): 12 | assert X.shape[1] == Y.shape[1] == self.d, (X.shape[1], Y.shape[1], self.d) 13 | coefficients = torch.tanh(self.A(weights)) 14 | res = torch.matmul(X * coefficients, Y.transpose(0, 1)) 15 | res = torch.nn.functional.softplus(res) - 0.5 16 | return res 17 | 18 | def forward(self, Xs, Ys, Ws): 19 | return [self._forward(X, Y, W) for X, Y, W in zip(Xs, Ys, Ws)] 20 | -------------------------------------------------------------------------------- /models/NGM/model_config.py: -------------------------------------------------------------------------------- 1 | from easydict import EasyDict as edict 2 | 3 | __C = edict() 4 | 5 | model_cfg = __C 6 | 7 | # NGM model options 8 | __C.NGM = edict() 9 | __C.NGM.FEATURE_CHANNEL = 512 10 | __C.NGM.SK_ITER_NUM = 10 11 | __C.NGM.SK_EPSILON = 1e-10 12 | __C.NGM.SK_TAU = 0.005 13 | __C.NGM.MGM_SK_TAU = 0.005 14 | __C.NGM.GNN_FEAT = [16, 16, 16] 15 | __C.NGM.GNN_LAYER = 3 16 | __C.NGM.GAUSSIAN_SIGMA = 1. 17 | __C.NGM.SIGMA3 = 1. 18 | __C.NGM.WEIGHT2 = 1. 19 | __C.NGM.WEIGHT3 = 1. 20 | __C.NGM.EDGE_FEATURE = 'cat' # 'cat' or 'geo' 21 | __C.NGM.ORDER3_FEATURE = 'none' # 'cat' or 'geo' or 'none' 22 | __C.NGM.FIRST_ORDER = True 23 | __C.NGM.EDGE_EMB = False 24 | __C.NGM.SK_EMB = 1 25 | __C.NGM.GUMBEL_SK = 0 # 0 for no gumbel, other wise for number of gumbel samples 26 | __C.NGM.UNIV_SIZE = -1 27 | __C.NGM.POSITIVE_EDGES = True 28 | -------------------------------------------------------------------------------- /src/utils/model_sl.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.nn import DataParallel 3 | 4 | 5 | def save_model(model, path): 6 | if isinstance(model, DataParallel): 7 | model = model.module 8 | 9 | torch.save(model.state_dict(), path) 10 | 11 | 12 | def load_model(model, path, strict=True): 13 | if isinstance(model, DataParallel): 14 | module = model.module 15 | else: 16 | module = model 17 | missing_keys, unexpected_keys = module.load_state_dict(torch.load(path), strict=strict) 18 | if len(unexpected_keys) > 0: 19 | print('Warning: Unexpected key(s) in state_dict: {}. '.format( 20 | ', '.join('"{}"'.format(k) for k in unexpected_keys))) 21 | if len(missing_keys) > 0: 22 | print('Warning: Missing key(s) in state_dict: {}. '.format( 23 | ', '.join('"{}"'.format(k) for k in missing_keys))) 24 | -------------------------------------------------------------------------------- /src/utils/pad_tensor.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | import torch.nn.functional as functional 4 | 5 | def pad_tensor(inp): 6 | """ 7 | Pad a list of input tensors into a list of tensors with same dimension 8 | :param inp: input tensor list 9 | :return: output tensor list 10 | """ 11 | assert type(inp[0]) == torch.Tensor 12 | it = iter(inp) 13 | t = next(it) 14 | max_shape = list(t.shape) 15 | while True: 16 | try: 17 | t = next(it) 18 | for i in range(len(max_shape)): 19 | max_shape[i] = int(max(max_shape[i], t.shape[i])) 20 | except StopIteration: 21 | break 22 | max_shape = np.array(max_shape) 23 | 24 | padded_ts = [] 25 | for t in inp: 26 | pad_pattern = np.zeros(2 * len(max_shape), dtype=np.int64) 27 | pad_pattern[::-2] = max_shape - np.array(t.shape) 28 | pad_pattern = tuple(pad_pattern.tolist()) 29 | padded_ts.append(functional.pad(t, pad_pattern, 'constant', 0)) 30 | 31 | return padded_ts -------------------------------------------------------------------------------- /src/utils/dup_stdout_manager.py: -------------------------------------------------------------------------------- 1 | import sys 2 | 3 | 4 | class DupStdoutFileWriter(object): 5 | def __init__(self, stdout, path, mode): 6 | self.path = path 7 | self._content = '' 8 | self._stdout = stdout 9 | self._file = open(path, mode) 10 | 11 | def write(self, msg): 12 | while '\n' in msg: 13 | pos = msg.find('\n') 14 | self._content += msg[:pos + 1] 15 | self.flush() 16 | msg = msg[pos + 1:] 17 | self._content += msg 18 | if len(self._content) > 1000: 19 | self.flush() 20 | 21 | def flush(self): 22 | self._stdout.write(self._content) 23 | self._stdout.flush() 24 | self._file.write(self._content) 25 | self._file.flush() 26 | self._content = '' 27 | 28 | def __del__(self): 29 | self._file.close() 30 | 31 | 32 | class DupStdoutFileManager(object): 33 | def __init__(self, path, mode='w+'): 34 | self.path = path 35 | self.mode = mode 36 | 37 | def __enter__(self): 38 | self._stdout = sys.stdout 39 | self._file = DupStdoutFileWriter(self._stdout, self.path, self.mode) 40 | sys.stdout = self._file 41 | 42 | def __exit__(self, exc_type, exc_value, traceback): 43 | sys.stdout = self._stdout -------------------------------------------------------------------------------- /models/NGM/geo_edge_feature.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import Tensor 3 | 4 | 5 | def geo_edge_feature(P: Tensor, G: Tensor, H: Tensor, norm_d=256, device=None): 6 | """ 7 | Compute geometric edge features [d, cos(theta), sin(theta)] 8 | Adjacency matrix is formed by A = G * H^T 9 | :param P: point set (b x num_nodes x 2) 10 | :param G: factorized graph partition G (b x num_nodes x num_edges) 11 | :param H: factorized graph partition H (b x num_nodes x num_edges) 12 | :param norm_d: normalize Euclidean distance by norm_d 13 | :param device: device 14 | :return: feature tensor (b x 3 x num_edges) 15 | """ 16 | if device is None: 17 | device = P.device 18 | 19 | p1 = torch.sum(torch.mul(P.unsqueeze(-2), G.unsqueeze(-1)), dim=1) # (b x num_edges x dim) 20 | p2 = torch.sum(torch.mul(P.unsqueeze(-2), H.unsqueeze(-1)), dim=1) 21 | 22 | d = torch.norm((p1 - p2) / (norm_d * torch.sum(G, dim=1)).unsqueeze(-1), dim=-1) # (b x num_edges) 23 | # non-existing elements are nan 24 | 25 | cos_theta = (p1[:, :, 0] - p2[:, :, 0]) / (d * norm_d) # non-existing elements are nan 26 | sin_theta = (p1[:, :, 1] - p2[:, :, 1]) / (d * norm_d) 27 | 28 | return torch.stack((d, cos_theta, sin_theta), dim=1).to(device) 29 | -------------------------------------------------------------------------------- /src/utils/print_easydict.py: -------------------------------------------------------------------------------- 1 | from easydict import EasyDict as edict 2 | 3 | def static_vars(**kwargs): 4 | def decorate(func): 5 | for k in kwargs: 6 | setattr(func, k, kwargs[k]) 7 | return func 8 | return decorate 9 | 10 | @static_vars(indent_cnt=0) 11 | def print_easydict(inp_dict: edict): 12 | for key, value in inp_dict.items(): 13 | if type(value) is edict or type(value) is dict: 14 | print('{}{}:'.format(' ' * 2 * print_easydict.indent_cnt, key)) 15 | print_easydict.indent_cnt += 1 16 | print_easydict(value) 17 | print_easydict.indent_cnt -= 1 18 | 19 | else: 20 | print('{}{}: {}'.format(' ' * 2 * print_easydict.indent_cnt, key, value)) 21 | 22 | @static_vars(indent_cnt=0) 23 | def print_easydict_str(inp_dict: edict): 24 | ret_str = '' 25 | for key, value in inp_dict.items(): 26 | if type(value) is edict or type(value) is dict: 27 | ret_str += '{}{}:\n'.format(' ' * 2 * print_easydict_str.indent_cnt, key) 28 | print_easydict_str.indent_cnt += 1 29 | ret_str += print_easydict_str(value) 30 | print_easydict_str.indent_cnt -= 1 31 | 32 | else: 33 | ret_str += '{}{}: {}\n'.format(' ' * 2 * print_easydict_str.indent_cnt, key, value) 34 | 35 | return ret_str 36 | -------------------------------------------------------------------------------- /src/rules.py: -------------------------------------------------------------------------------- 1 | #rule_dict = { 2 | #'SeatBase': ['B_WheelIntersection', 'F_WheelIntersection', 'CranksetCenter', 'HandleCenter'], 3 | #'CranksetCenter': ['B_WheelIntersection', 'F_WheelIntersection', 'SeatBase'], 4 | #'HandleCenter': ['L_HandleTip', 'R_HandleTip', 'SeatBase'], 5 | #'L_HandleTip': ['HandleCenter'], 6 | #'R_HandleTip': ['HandleCenter'], 7 | #'B_WheelIntersection': ['B_WheelCenter', 'B_WheelEnd', 'CranksetCenter'], 8 | #'B_WheelEnd': ['B_WheelIntersection', 'B_WheelCenter'], 9 | #'B_WheelCenter': ['B_WheelIntersection', 'B_WheelEnd'], 10 | #'F_WheelIntersection': ['F_WheelCenter', 'F_WheelEnd', 'CranksetCenter', 'HandleCenter'], 11 | #'F_WheelEnd': ['F_WheelIntersection', 'F_WheelCenter'], 12 | #'F_WheelCenter': ['F_WheelIntersection', 'F_WheelEnd'], 13 | #} 14 | 15 | rule_dict = { 16 | 'BackRest_Top_Left': ['BackRest_Top_Right', 'Seat_Left_Back'], 17 | 'BackRest_Top_Right': ['BackRest_Top_Left', 'Seat_Right_Back'], 18 | 'Leg_Left_Back': ['Seat_Left_Back'], 19 | 'Leg_Left_Front': ['Seat_Left_Front'], 20 | 'Leg_Right_Back': ['Seat_Right_Back'], 21 | 'Leg_Right_Front': ['Seat_Right_Front'], 22 | 'Seat_Left_Back': ['Leg_Left_Back', 'BackRest_Top_Left', 'Seat_Left_Front', 'Seat_Right_Back'], 23 | 'Seat_Left_Front': ['Leg_Left_Front', 'Seat_Left_Back', 'Seat_Right_Front', ], 24 | 'Seat_Right_Back': ['Leg_Right_Back', 'BackRest_Top_Right', 'Seat_Right_Front', 'Seat_Left_Back'], 25 | 'Seat_Right_Front': ['Leg_Right_Front', 'Seat_Right_Back', 'Seat_Left_Front'] 26 | } -------------------------------------------------------------------------------- /models/GMN/voting_layer.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | 5 | class Voting(nn.Module): 6 | """ 7 | Voting Layer computes a new row-stotatic matrix with softmax. A large number (alpha) is multiplied to the input 8 | stochastic matrix to scale up the difference. 9 | Parameter: value multiplied before softmax alpha 10 | threshold that will ignore such points while calculating displacement in pixels pixel_thresh 11 | Input: permutation or doubly stochastic matrix s 12 | ///point set on source image P_src 13 | ///point set on target image P_tgt 14 | ground truth number of effective points in source image ns_gt 15 | Output: softmax matrix s 16 | """ 17 | def __init__(self, alpha=200, pixel_thresh=None): 18 | super(Voting, self).__init__() 19 | self.alpha = alpha 20 | self.softmax = nn.Softmax(dim=-1) # Voting among columns 21 | self.pixel_thresh = pixel_thresh 22 | 23 | def forward(self, s, nrow_gt, ncol_gt=None): 24 | # TODO discard dummy nodes & far away nodes 25 | ret_s = torch.zeros_like(s) 26 | # filter dummy nodes 27 | for b, n in enumerate(nrow_gt): 28 | if ncol_gt is None: 29 | ret_s[b, 0:n, :] = \ 30 | self.softmax(self.alpha * s[b, 0:n, :]) 31 | else: 32 | ret_s[b, 0:n, 0:ncol_gt[b]] =\ 33 | self.softmax(self.alpha * s[b, 0:n, 0:ncol_gt[b]]) 34 | 35 | return ret_s 36 | -------------------------------------------------------------------------------- /models/GMN/displacement_layer.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | 5 | class Displacement(nn.Module): 6 | """ 7 | Displacement Layer computes the displacement vector for each point in the source image, with its corresponding point 8 | (or points) in target image. The output is a displacement matrix constructed from all displacement vectors. 9 | This metric measures the shift from source point to predicted target point, and can be applied for matching 10 | accuracy. 11 | Together with displacement matrix d, this function will also return a grad_mask, which helps to filter out dummy 12 | nodes in practice. 13 | d = s * P_tgt - P_src 14 | Proposed by Zanfir et al. Deep Learning of Graph Matching. CVPR 2018. 15 | Input: permutation or doubly stochastic matrix s 16 | point set on source image P_src 17 | point set on target image P_tgt 18 | (optional) ground truth number of effective points in source image ns_gt 19 | Output: displacement matrix d 20 | mask for dummy nodes grad_mask. If ns_gt=None, it will not be calculated and None is returned. 21 | """ 22 | def __init__(self): 23 | super(Displacement, self).__init__() 24 | 25 | def forward(self, s, P_src, P_tgt, ns_gt=None): 26 | if ns_gt is None: 27 | max_n = s.shape[1] 28 | P_src = P_src[:, 0:max_n, :] 29 | grad_mask = None 30 | else: 31 | grad_mask = torch.zeros_like(P_src) 32 | for b, n in enumerate(ns_gt): 33 | grad_mask[b, 0:n] = 1 34 | 35 | d = torch.matmul(s, P_tgt) - P_src 36 | return d, grad_mask -------------------------------------------------------------------------------- /src/qap_solvers/spectral_matching.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from src.utils.sparse import sbmm 4 | 5 | 6 | class SpectralMatching(nn.Module): 7 | """ 8 | Spectral Graph Matching solver. 9 | Also known as Power Iteration layer, which computes the leading eigenvector of input matrix. 10 | For every iteration, 11 | v_k+1 = M * v_k / ||M * v_k||_2 12 | Parameter: maximum iteration max_iter 13 | Input: input matrix M 14 | (optional) initialization vector v0. If not specified, v0 will be initialized with all 1. 15 | Output: computed eigenvector v 16 | """ 17 | def __init__(self, max_iter=50, stop_thresh=2e-7): 18 | super(SpectralMatching, self).__init__() 19 | self.max_iter = max_iter 20 | self.stop_thresh = stop_thresh 21 | 22 | def forward(self, M, v0=None, **kwargs): 23 | batch_num = M.shape[0] 24 | mn = M.shape[1] 25 | if v0 is None: 26 | v0 = torch.ones(batch_num, mn, 1, dtype=M.dtype, device=M.device) 27 | 28 | v = vlast = v0 29 | for i in range(self.max_iter): 30 | if M.is_sparse: 31 | v = sbmm(M, v) 32 | else: 33 | v = torch.bmm(M, v) 34 | n = torch.norm(v, p=2, dim=1) 35 | v = torch.matmul(v, (1 / n).view(batch_num, 1, 1)) 36 | if torch.norm(v - vlast) < self.stop_thresh: 37 | return v.view(batch_num, -1) 38 | vlast = v 39 | 40 | return v.view(batch_num, -1) 41 | 42 | 43 | if __name__ == '__main__': 44 | from torch.autograd import gradcheck 45 | input = (torch.randn(3, 40, 40, dtype=torch.double, requires_grad=True),) 46 | 47 | pi = SpectralMatching() 48 | 49 | test = gradcheck(pi, input, eps=1e-6, atol=1e-4) 50 | print(test) -------------------------------------------------------------------------------- /experiments/eval.yaml: -------------------------------------------------------------------------------- 1 | MODEL_NAME: vgg16_ngmv2 2 | DATASET_NAME: voc 3 | 4 | DATASET_FULL_NAME: PascalVOC 5 | 6 | MODULE: models.NGM.model_v2 7 | 8 | BACKBONE: VGG16_bn_final 9 | 10 | BATCH_SIZE: 8 11 | DATALOADER_NUM: 2 12 | FP16: False 13 | 14 | RANDOM_SEED: 123 15 | 16 | PRETRAINED_PATH: "/home/baoqingquan/RobustMatch/pretrained/new_config1.pt" # path that needs to be specified to load model parameters 17 | 18 | # available GPU ids 19 | GPUS: 20 | - 0 21 | 22 | # Problem configuration 23 | PROBLEM: 24 | TYPE: 2GM 25 | RESCALE: # rescaled image size 26 | - 256 27 | - 256 28 | 29 | # Graph construction settings 30 | GRAPH: 31 | SRC_GRAPH_CONSTRUCT: tri 32 | TGT_GRAPH_CONSTRUCT: tri 33 | SYM_ADJACENCY: True 34 | 35 | # model parameters 36 | NGM: 37 | EDGE_FEATURE: cat 38 | FEATURE_CHANNEL: 512 39 | SK_ITER_NUM: 20 40 | SK_EPSILON: 1.0e-10 41 | SK_TAU: 0.05 42 | GNN_FEAT: 43 | - 16 44 | - 16 45 | - 16 46 | GNN_LAYER: 3 47 | GAUSSIAN_SIGMA: 1. 48 | SK_EMB: 1 49 | FIRST_ORDER: True 50 | EDGE_EMB: False 51 | 52 | TRAIN: 53 | MODE: eval 54 | 55 | # Evaluation settings 56 | EVAL: 57 | EPOCH: 15 # where to continue evlauation based on training epochs 58 | SAMPLES: 200 # number of tested pairs for each class 59 | NUM_EPOCH: 4 # evaluation period during training 60 | MODE: single # choices=['clean', 'single', 'all'] 61 | 62 | # Attack params for AT 63 | ATTACK: 64 | # basic params 65 | EPSILON_FEATURE: 8 66 | EPSILON_LOCALITY: 8 67 | EVAL_ALPHA: 0.25 68 | EVAL_STEP: 10 69 | RESTARTS: 1 70 | EARLY_STOP_RATIO: 0. 71 | # attack attributes 72 | TYPE: pgd # attack optimization way, choices = ['pgd', 'none', 'momentum', 'random'] 73 | OBJ_TYPE: pixel # attack input object, choices = ['pixel', 'pos', 'struc', 'pos+struc', 'pixel+pos', 'pixel+pos+struc'] 74 | LOSS_FUNC: perm # attack loss objective, choices = ['perm', 'ourloss', 'hamming', 'cw', ...] -------------------------------------------------------------------------------- /src/AWP/awp.py: -------------------------------------------------------------------------------- 1 | # This is a simple implementation of AWP for Standard Adversarial Training (Madry) 2 | import copy 3 | import torch.nn as nn 4 | import torch.optim as optim 5 | import torch 6 | EPS = 1E-20 7 | 8 | 9 | def normalize(perturbations, weights): 10 | perturbations.mul_(weights.norm()/(perturbations.norm() + EPS)) 11 | 12 | 13 | def normalize_grad_by_weights(weights, ref_weights): 14 | for w, ref_w in zip(weights, ref_weights): 15 | if w.dim() <= 1: 16 | w.grad.data.fill_(0) # ignore perturbations with 1 dimension (e.g. BN, bias) 17 | else: 18 | normalize(w.grad.data, ref_w) 19 | 20 | 21 | class AdvWeightPerturb(object): 22 | """ 23 | This is an implementation of AWP ONLY for Standard adversarial training 24 | """ 25 | def __init__(self, model, eta, nb_iter=1): 26 | super(AdvWeightPerturb, self).__init__() 27 | self.eta = eta 28 | self.nb_iter = nb_iter 29 | self.model = model 30 | self.optim = optim.SGD(model.parameters(), lr=eta/nb_iter) 31 | self.criterion = nn.CrossEntropyLoss() 32 | self.diff = None 33 | 34 | def perturb(self, X_adv, y): 35 | # store the original weight 36 | old_w = copy.deepcopy([p.data for p in self.model.parameters()]) 37 | 38 | # perturb the model 39 | for idx in range(self.nb_iter): 40 | self.optim.zero_grad() 41 | outputs = self.model(X_adv) 42 | loss = - self.criterion(outputs, y) 43 | loss.backward() 44 | 45 | # normalize the gradient 46 | normalize_grad_by_weights(self.model.parameters(), old_w) 47 | 48 | self.optim.step() 49 | 50 | # calculate the weight perturbation 51 | self.diff = [w1 - w2 for w1, w2 in zip(self.model.parameters(), old_w)] 52 | 53 | def restore(self): 54 | for w, v in zip(self.model.parameters(), self.diff): 55 | w.data.sub_(v.data) 56 | -------------------------------------------------------------------------------- /src/extension/sparse_dot/csr_dot_diag_cuda.cu: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | #include 4 | #include 5 | 6 | #include 7 | 8 | 9 | template 10 | __global__ void csr_dot_diag_cuda_kernel( 11 | const int64_t* __restrict__ t1_indices, 12 | const int64_t* __restrict__ t1_indptr, 13 | const scalar_t* __restrict__ t1_data, 14 | const scalar_t* __restrict__ t2, 15 | scalar_t* __restrict__ outp_data, 16 | const int64_t out_h, 17 | const int64_t out_w 18 | ) 19 | { 20 | const int64_t i = blockIdx.x * blockDim.x + threadIdx.x; 21 | const int64_t b = blockIdx.y; 22 | 23 | if (i < out_h) 24 | { 25 | const int64_t start = t1_indptr[b * out_h + i]; 26 | const int64_t stop = t1_indptr[b * out_h + i + 1]; 27 | 28 | for (int64_t data_idx = start; data_idx < stop; data_idx++) 29 | { 30 | int64_t row_idx = t1_indices[data_idx]; 31 | outp_data[data_idx] = t1_data[data_idx] * t2[b * out_w + row_idx]; 32 | } 33 | } 34 | } 35 | 36 | 37 | std::vector csr_dot_diag_cuda( 38 | at::Tensor t1_indices, 39 | at::Tensor t1_indptr, 40 | at::Tensor t1_data, 41 | at::Tensor t2, 42 | int64_t batch_size, 43 | int64_t out_h, 44 | int64_t out_w 45 | ){ 46 | auto outp_indices = at::clone(t1_indices); 47 | auto outp_indptr = at::clone(t1_indptr); 48 | auto outp_data = at::zeros_like(t1_data); 49 | 50 | const int threads = 1024; 51 | const dim3 blocks((out_h + threads - 1) / threads, batch_size); 52 | 53 | AT_DISPATCH_FLOATING_TYPES_AND_HALF(t1_data.type(), "csr_dot_diag_cuda", ([&] { 54 | csr_dot_diag_cuda_kernel<<>>( 55 | t1_indices.data(), 56 | t1_indptr.data(), 57 | t1_data.data(), 58 | t2.data(), 59 | outp_data.data(), 60 | out_h, 61 | out_w); 62 | })); 63 | 64 | return {outp_indices, outp_indptr, outp_data}; 65 | } 66 | -------------------------------------------------------------------------------- /src/utils/utils_awp.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from collections import OrderedDict 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | from src.utils.config import cfg 6 | EPS = 1E-20 7 | 8 | def diff_in_weights(model, proxy): 9 | diff_dict = OrderedDict() 10 | model_state_dict = model.state_dict() 11 | proxy_state_dict = proxy.state_dict() 12 | for (old_k, old_w), (new_k, new_w) in zip(model_state_dict.items(), proxy_state_dict.items()): 13 | if len(old_w.size()) <= 1: 14 | continue 15 | if 'weight' in old_k: 16 | diff_w = new_w - old_w 17 | diff_dict[old_k] = old_w.norm() / (diff_w.norm() + EPS) * diff_w 18 | return diff_dict 19 | 20 | 21 | def add_into_weights(model, diff, coeff=1.0): 22 | names_in_diff = diff.keys() 23 | with torch.no_grad(): 24 | for name, param in model.named_parameters(): 25 | if name in names_in_diff: 26 | param.add_(coeff * diff[name]) 27 | 28 | 29 | class AdvWeightPerturb(object): 30 | def __init__(self, model, proxy, proxy_optim, gamma): 31 | super(AdvWeightPerturb, self).__init__() 32 | self.model = model 33 | self.proxy = proxy 34 | self.proxy_optim = proxy_optim 35 | self.gamma = gamma 36 | 37 | def calc_awp(self, inputs, criterion, inputs_att=None): 38 | self.proxy.load_state_dict(self.model.state_dict()) 39 | self.proxy.train() 40 | 41 | if inputs_att is not None: 42 | with torch.no_grad(): 43 | outputs_att = self.proxy(inputs_att) 44 | loss = -1. * criterion(self.proxy(inputs), outputs_att) 45 | else: 46 | loss = -1. * criterion(self.proxy(inputs)) 47 | self.proxy_optim.zero_grad() 48 | loss.backward() 49 | self.proxy_optim.step() 50 | 51 | # the adversary weight perturb 52 | diff = diff_in_weights(self.model, self.proxy) 53 | return diff 54 | 55 | def perturb(self, diff): 56 | add_into_weights(self.model, diff, coeff=1.0 * self.gamma) 57 | 58 | def restore(self, diff): 59 | add_into_weights(self.model, diff, coeff=-1.0 * self.gamma) 60 | -------------------------------------------------------------------------------- /src/qap_solvers/rrwm.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | from src.lap_solvers.sinkhorn import Sinkhorn as Sinkhorn 5 | 6 | 7 | class RRWM(nn.Module): 8 | """ 9 | RRWM solver for graph matching (QAP), implemented by power iteration with Sinkhorn reweighted jumps. 10 | Parameter: maximum iteration max_iter 11 | Input: input matrix M 12 | maximum size of source graph num_src 13 | sizes of source graph in batch ns_src 14 | sizes of target graph in batch ns_tgt 15 | (optional) initialization vector v0. If not specified, v0 will be initialized with all 1. 16 | Output: computed eigenvector v 17 | """ 18 | def __init__(self, max_iter=50, sk_iter=20, alpha=0.2, beta=30): 19 | super(RRWM, self).__init__() 20 | self.max_iter = max_iter 21 | self.alpha = alpha 22 | self.beta = beta 23 | self.sk = Sinkhorn(max_iter=sk_iter,log_forward=False) 24 | 25 | def forward(self, M, num_src, ns_src, ns_tgt, v0=None): 26 | d = M.sum(dim=2, keepdim=True) 27 | dmax = d.max(dim=1, keepdim=True).values 28 | M = M / (dmax + d.min() * 1e-5) 29 | 30 | batch_num = M.shape[0] 31 | mn = M.shape[1] 32 | if v0 is None: 33 | v0 = torch.zeros(batch_num, num_src, mn // num_src, dtype=M.dtype, device=M.device) 34 | for b in range(batch_num): 35 | v0[b, 0:ns_src[b], 0:ns_tgt[b]] = torch.tensor(1.) / (ns_src[b] * ns_tgt[b]) 36 | 37 | v0 = v0.transpose(1, 2).reshape(batch_num, mn, 1) 38 | 39 | v = v0 40 | for i in range(self.max_iter): 41 | v = torch.bmm(M, v) 42 | last_v = v 43 | n = torch.norm(v, p=1, dim=1, keepdim=True) 44 | v = v / n 45 | s = v.view(batch_num, -1, num_src).transpose(1, 2) 46 | s = torch.exp(self.beta * s / s.max(dim=1, keepdim=True).values.max(dim=2, keepdim=True).values) 47 | 48 | v = self.alpha * self.sk(s, ns_src, ns_tgt).transpose(1, 2).reshape(batch_num, mn, 1) + (1 - self.alpha) * v 49 | n = torch.norm(v, p=1, dim=1, keepdim=True) 50 | v = torch.matmul(v, 1 / n) 51 | 52 | if torch.norm(v - last_v) < 1e-5: 53 | break 54 | 55 | return v.view(batch_num, -1) 56 | -------------------------------------------------------------------------------- /experiments/eval_blackbox.yaml: -------------------------------------------------------------------------------- 1 | MODEL_NAME: vgg16_ngmv2 2 | DATASET_NAME: voc 3 | 4 | DATASET_FULL_NAME: PascalVOC 5 | 6 | MODULE: models.NGM.model_v2 7 | 8 | BACKBONE: VGG16_bn_final 9 | 10 | BATCH_SIZE: 8 11 | DATALOADER_NUM: 2 12 | FP16: False 13 | 14 | RANDOM_SEED: 123 15 | 16 | PRETRAINED_PATH: "" # path that needs to be specified to load model parameters 17 | 18 | VICTIM_MODEL_NAME: vgg16_ngmv2 19 | VICTIM_MODULE: models.NGM.model_v2 20 | VICTIM_PATH: "" # path that needs to be specified to load model parameters 21 | 22 | 23 | # available GPU ids 24 | GPUS: 25 | - 0 26 | 27 | # Problem configuration 28 | PROBLEM: 29 | TYPE: 2GM 30 | RESCALE: # rescaled image size 31 | - 256 32 | - 256 33 | 34 | # Graph construction settings 35 | GRAPH: 36 | SRC_GRAPH_CONSTRUCT: tri 37 | TGT_GRAPH_CONSTRUCT: tri 38 | SYM_ADJACENCY: True 39 | 40 | # model parameters 41 | NGM: 42 | EDGE_FEATURE: cat 43 | FEATURE_CHANNEL: 512 44 | SK_ITER_NUM: 20 45 | SK_EPSILON: 1.0e-10 46 | SK_TAU: 0.05 47 | GNN_FEAT: 48 | - 16 49 | - 16 50 | - 16 51 | GNN_LAYER: 3 52 | GAUSSIAN_SIGMA: 1. 53 | SK_EMB: 1 54 | FIRST_ORDER: True 55 | EDGE_EMB: False 56 | 57 | # If you need to attack other model, e.g., BBGM, 58 | # you should also define the following config: 59 | #BBGM: 60 | # FEATURE_CHANNEL: 1024 61 | # SOLVER_NAME: LPMP 62 | # LAMBDA_VAL: 80.0 63 | # SOLVER_PARAMS: 64 | # timeout: 1000 65 | # primalComputationInterval: 10 66 | # maxIter: 10 67 | 68 | TRAIN: 69 | MODE: eval 70 | 71 | # Evaluation settings 72 | EVAL: 73 | EPOCH: 15 # where to continue evlauation based on training epochs 74 | SAMPLES: 1000 # number of tested pairs for each class 75 | NUM_EPOCH: 4 # evaluation period during training 76 | MODE: all # choices=['clean', 'single', 'all'] 77 | 78 | # Attack params for AT 79 | ATTACK: 80 | # basic params 81 | EPSILON_FEATURE: 8 82 | EPSILON_LOCALITY: 8 83 | EVAL_ALPHA: 0.25 84 | EVAL_STEP: 10 85 | RESTARTS: 1 86 | EARLY_STOP_RATIO: 0. 87 | # attack attributes 88 | TYPE: pgd # attack optimization way, choices = ['pgd', 'none', 'momentum', 'random'] 89 | OBJ_TYPE: pixel+pos+struc # attack input object, choices = ['pixel', 'pos', 'struc', 'pos+struc', 'pixel+pos', 'pixel+pos+struc'] 90 | LOSS_FUNC: perm # attack loss objective, choices = ['perm', 'ourloss', 'hamming', 'cw', ...] -------------------------------------------------------------------------------- /src/displacement_layer.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from torch import Tensor 4 | 5 | 6 | class Displacement(nn.Module): 7 | r""" 8 | Displacement Layer computes the displacement vector for each point in the source image, with its corresponding point 9 | (or points) in target image. 10 | 11 | The output is a displacement matrix constructed from all displacement vectors. 12 | This metric measures the shift from source point to predicted target point, and can be applied for matching 13 | accuracy. 14 | 15 | Together with displacement matrix d, this function will also return a grad_mask, which helps to filter out dummy 16 | nodes in practice. 17 | 18 | .. math:: 19 | \mathbf{d}_i = \sum_{j \in V_2} \left( \mathbf{S}_{i, j} P_{2j} \right)- P_{1i} 20 | 21 | Proposed by `"Zanfir et al. Deep Learning of Graph Matching. CVPR 2018." 22 | `_ 23 | """ 24 | def __init__(self): 25 | super(Displacement, self).__init__() 26 | 27 | def forward(self, s: Tensor, P_src: Tensor, P_tgt: Tensor, ns_gt: Tensor=None): 28 | r""" 29 | :param s: :math:`(b\times n_1 \times n_2)` permutation or doubly stochastic matrix. :math:`b`: batch size. 30 | :math:`n_1`: number of nodes in source image. :math:`n_2`: number of nodes in target image 31 | :param P_src: :math:`(b\times n_1 \times 2)` point set on source image 32 | :param P_tgt: :math:`(b\times n_2 \times 2)` point set on target image 33 | :param ns_gt: :math:`(b)` number of exact pairs. We support batched instances with different number of nodes, 34 | therefore ``ns_gt`` is required to specify the exact number of nodes of each instance in the batch. 35 | :return: displacement matrix d, 36 | mask for dummy nodes grad_mask. If ``ns_gt=None``, it will not be calculated and None is returned. 37 | """ 38 | if ns_gt is None: 39 | max_n = s.shape[1] 40 | P_src = P_src[:, 0:max_n, :] 41 | grad_mask = None 42 | else: 43 | grad_mask = torch.zeros_like(P_src) 44 | for b, n in enumerate(ns_gt): 45 | grad_mask[b, 0:n] = 1 46 | 47 | d = torch.matmul(s, P_tgt) - P_src 48 | return d, grad_mask -------------------------------------------------------------------------------- /src/lap_solvers/hungarian.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import scipy.optimize as opt 3 | import numpy as np 4 | from multiprocessing import Pool 5 | from torch import Tensor 6 | 7 | 8 | def hungarian(s: Tensor, n1: Tensor=None, n2: Tensor=None, nproc: int=1) -> Tensor: 9 | r""" 10 | Solve optimal LAP permutation by hungarian algorithm. The time cost is :math:`O(n^3)`. 11 | 12 | :param s: :math:`(b\times n_1 \times n_2)` input 3d tensor. :math:`b`: batch size 13 | :param n1: :math:`(b)` number of objects in dim1 14 | :param n2: :math:`(b)` number of objects in dim2 15 | :param nproc: number of parallel processes (default: ``nproc=1`` for no parallel) 16 | :return: :math:`(b\times n_1 \times n_2)` optimal permutation matrix 17 | 18 | .. note:: 19 | We support batched instances with different number of nodes, therefore ``n1`` and ``n2`` are 20 | required to specify the exact number of objects of each dimension in the batch. If not specified, we assume 21 | the batched matrices are not padded. 22 | """ 23 | if len(s.shape) == 2: 24 | s = s.unsqueeze(0) 25 | matrix_input = True 26 | elif len(s.shape) == 3: 27 | matrix_input = False 28 | else: 29 | raise ValueError('input data shape not understood: {}'.format(s.shape)) 30 | 31 | device = s.device 32 | batch_num = s.shape[0] 33 | 34 | perm_mat = s.cpu().detach().numpy() * -1 35 | if n1 is not None: 36 | n1 = n1.cpu().numpy() 37 | else: 38 | n1 = [None] * batch_num 39 | if n2 is not None: 40 | n2 = n2.cpu().numpy() 41 | else: 42 | n2 = [None] * batch_num 43 | 44 | if nproc > 1: 45 | with Pool(processes=nproc) as pool: 46 | mapresult = pool.starmap_async(_hung_kernel, zip(perm_mat, n1, n2)) 47 | perm_mat = np.stack(mapresult.get()) 48 | else: 49 | perm_mat = np.stack([_hung_kernel(perm_mat[b], n1[b], n2[b]) for b in range(batch_num)]) 50 | 51 | perm_mat = torch.from_numpy(perm_mat).to(device) 52 | 53 | if matrix_input: 54 | perm_mat.squeeze_(0) 55 | 56 | return perm_mat 57 | 58 | def _hung_kernel(s: torch.Tensor, n1=None, n2=None): 59 | if n1 is None: 60 | n1 = s.shape[0] 61 | if n2 is None: 62 | n2 = s.shape[1] 63 | row, col = opt.linear_sum_assignment(s[:n1, :n2]) 64 | perm_mat = np.zeros_like(s) 65 | perm_mat[row, col] = 1 66 | return perm_mat -------------------------------------------------------------------------------- /src/extension/bilinear_diag/bilinear_diag_cuda.cu: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | #include 4 | #include 5 | 6 | 7 | template 8 | __global__ void bilinear_diag_csc_cuda_kernel( 9 | const int64_t* __restrict__ t1_indices, 10 | const int64_t* __restrict__ t1_indptr, 11 | const scalar_t* __restrict__ t1_data, 12 | const scalar_t* __restrict__ t2, 13 | const int64_t* __restrict__ t3_indices, 14 | const int64_t* __restrict__ t3_indptr, 15 | const scalar_t* __restrict__ t3_data, 16 | scalar_t* __restrict__ outp, 17 | const int64_t xlen, 18 | const int64_t feat_size 19 | ) 20 | { 21 | const int64_t i = blockIdx.x * blockDim.x + threadIdx.x; 22 | const int64_t b = blockIdx.y; 23 | 24 | if (i < xlen) 25 | { 26 | const int64_t ptr_idx = b * xlen + i; 27 | const int64_t t1_start = t1_indptr[ptr_idx]; 28 | const int64_t t1_stop = t1_indptr[ptr_idx + 1]; 29 | const int64_t t3_start = t3_indptr[ptr_idx]; 30 | const int64_t t3_stop = t3_indptr[ptr_idx + 1]; 31 | 32 | scalar_t _outp = 0; 33 | 34 | for (int64_t t1_idx = t1_start; t1_idx < t1_stop; t1_idx++) 35 | { 36 | for (int64_t t3_idx = t3_start; t3_idx < t3_stop; t3_idx++) 37 | { 38 | _outp += t2[b * feat_size * feat_size + t1_indices[t1_idx] * feat_size + t3_indices[t3_idx]] 39 | * t1_data[t1_idx] * t3_data[t3_idx]; 40 | } 41 | } 42 | outp[b * xlen + i] = _outp; 43 | } 44 | } 45 | 46 | 47 | at::Tensor bilinear_diag_csc_cuda( 48 | at::Tensor t1_indices, 49 | at::Tensor t1_indptr, 50 | at::Tensor t1_data, 51 | at::Tensor t2, 52 | at::Tensor t3_indices, 53 | at::Tensor t3_indptr, 54 | at::Tensor t3_data, 55 | int64_t batch_size, 56 | int64_t xlen 57 | ){ 58 | auto outp = at::zeros({batch_size, xlen}, t2.type()); 59 | auto feat_size = t2.size(1); 60 | 61 | const int threads = 1024; 62 | const dim3 blocks((xlen + threads - 1) / threads, batch_size); 63 | 64 | AT_DISPATCH_FLOATING_TYPES_AND_HALF(t2.type(), "bilinear_diag_csc_cuda", ([&] { 65 | bilinear_diag_csc_cuda_kernel<<>>( 66 | t1_indices.data(), 67 | t1_indptr.data(), 68 | t1_data.data(), 69 | t2.data(), 70 | t3_indices.data(), 71 | t3_indptr.data(), 72 | t3_data.data(), 73 | outp.data(), 74 | xlen, 75 | feat_size); 76 | })); 77 | 78 | return outp; 79 | } 80 | -------------------------------------------------------------------------------- /src/qap_solvers/rrwhm.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | from src.lap_solvers.sinkhorn import Sinkhorn as Sinkhorn 5 | 6 | 7 | class RRWHM(nn.Module): 8 | """ 9 | RRWHM solver for hyper graph matching, implemented by tensor power iteration with Sinkhorn reweighted jumps. 10 | Parameter: maximum iteration max_iter 11 | Input: input tensor H 12 | maximum size of source graph num_src 13 | sizes of source graph in batch ns_src 14 | sizes of target graph in batch ns_tgt 15 | (optional) initialization vector v0. If not specified, v0 will be initialized with all 1. 16 | Output: computed eigenvector v 17 | """ 18 | def __init__(self, max_iter=50, sk_iter=20, alpha=0.2, beta=30): 19 | super(RRWHM, self).__init__() 20 | self.max_iter = max_iter 21 | self.alpha = alpha 22 | self.beta = beta 23 | self.sk = Sinkhorn(max_iter=sk_iter,log_forward=False) 24 | 25 | def forward(self, H, num_src, ns_src, ns_tgt, v0=None): 26 | order = len(H.shape) - 1 27 | sum_dims = [i+2 for i in range(order-1)] 28 | d = H.sum(dim=sum_dims, keepdim=True) 29 | dmax = d.max(dim=1, keepdim=True).values 30 | H = H / dmax 31 | 32 | batch_num = H.shape[0] 33 | mn = H.shape[1] 34 | if v0 is None: 35 | v0 = torch.zeros(batch_num, num_src, mn // num_src, dtype=H.dtype, device=H.device) 36 | for b in range(batch_num): 37 | v0[b, 0:ns_src[b], 0:ns_tgt[b]] = torch.tensor(1.) / (ns_src[b] * ns_tgt[b]) 38 | 39 | v0 = v0.transpose(1, 2).reshape(batch_num, mn, 1) 40 | 41 | v = v0 42 | for i in range(self.max_iter): 43 | H_red = H.unsqueeze(-1) 44 | for o in range(order - 1): 45 | v_shape = [v.shape[0]] + [1] * (order - 1 - o) + list(v.shape[1:]) 46 | H_red = torch.sum(torch.mul(H_red, v.view(*v_shape)), dim=-2) 47 | v = H_red 48 | last_v = v 49 | n = torch.norm(v, p=1, dim=1, keepdim=True) 50 | v = v / n 51 | s = v.view(batch_num, -1, num_src).transpose(1, 2) 52 | s = torch.exp(self.beta * s / s.max(dim=1, keepdim=True).values.max(dim=2, keepdim=True).values) 53 | 54 | v = self.alpha * self.sk(s, ns_src, ns_tgt).transpose(1, 2).reshape(batch_num, mn, 1) + (1 - self.alpha) * v 55 | n = torch.norm(v, p=1, dim=1, keepdim=True) 56 | v = torch.matmul(v, 1 / n) 57 | 58 | if torch.norm(v - last_v) < 1e-5: 59 | break 60 | 61 | return v.view(batch_num, -1) 62 | -------------------------------------------------------------------------------- /experiments/config1.yaml: -------------------------------------------------------------------------------- 1 | MODEL_NAME: vgg16_ngmv2 2 | DATASET_NAME: voc 3 | 4 | DATASET_FULL_NAME: PascalVOC 5 | 6 | MODULE: models.NGM.model_v2 7 | 8 | BACKBONE: VGG16_bn_final 9 | 10 | BATCH_SIZE: 8 11 | DATALOADER_NUM: 2 12 | FP16: False 13 | 14 | RANDOM_SEED: 123 15 | 16 | # available GPU ids 17 | GPUS: 18 | - 0 19 | 20 | # Problem configuration 21 | PROBLEM: 22 | TYPE: 2GM 23 | RESCALE: # rescaled image size 24 | - 256 25 | - 256 26 | 27 | # Graph construction settings 28 | GRAPH: 29 | SRC_GRAPH_CONSTRUCT: tri 30 | TGT_GRAPH_CONSTRUCT: tri 31 | SYM_ADJACENCY: True 32 | 33 | # Training settings 34 | TRAIN: 35 | # start, end epochs 36 | START_EPOCH: 0 37 | NUM_EPOCHS: 16 38 | 39 | MODE: 2step # choices = ['at', '2step'] 40 | LOSS_FUNC: ourloss # choices = ['perm', 'ourloss', 'hamming', ...] 41 | 42 | # for '2step' mode 43 | BURN_IN_PERIOD: 5 44 | 45 | # for AAR regularizer 46 | REG_LEVEL: 1 47 | REG_RATIO: 1.5 48 | 49 | # whether to synchronize optimization way for min-max AT framework, default: false 50 | SYNC_MINMAX: 0 51 | 52 | # set BN mode as evaluation during generating adversarial examples 53 | BN_EVAL: True 54 | 55 | OPTIMIZER: Adam 56 | # learning rate 57 | LR: 2.e-3 58 | SEPARATE_BACKBONE_LR: True 59 | BACKBONE_LR: 2.e-5 60 | MOMENTUM: 0.9 61 | LR_DECAY: 0.5 62 | LR_STEP: # (in epochs) 63 | - 2 64 | - 4 65 | - 6 66 | - 8 67 | - 10 68 | 69 | EPOCH_ITERS: 2000 # iterations per epoch 70 | 71 | CLASS: none 72 | 73 | # Evaluation settings 74 | EVAL: 75 | EPOCH: 15 # where to continue evlauation based on training epochs 76 | SAMPLES: 1000 # number of tested pairs for each class 77 | NUM_EPOCH: 4 # evaluation period during training 78 | 79 | # model parameters 80 | NGM: 81 | EDGE_FEATURE: cat 82 | FEATURE_CHANNEL: 512 83 | SK_ITER_NUM: 20 84 | SK_EPSILON: 1.0e-10 85 | SK_TAU: 0.05 86 | GNN_FEAT: 87 | - 16 88 | - 16 89 | - 16 90 | GNN_LAYER: 3 91 | GAUSSIAN_SIGMA: 1. 92 | SK_EMB: 1 93 | FIRST_ORDER: True 94 | EDGE_EMB: False 95 | 96 | # Attack params for AT 97 | ATTACK: 98 | # basic params 99 | EPSILON_FEATURE: 8 100 | EPSILON_LOCALITY: 8 101 | ALPHA: 1. 102 | STEP: 1 103 | EVAL_ALPHA: 0.25 104 | EVAL_STEP: 10 105 | RESTARTS: 1 106 | EARLY_STOP_RATIO: 0. 107 | # attack attributes 108 | TYPE: pgd # attack optimization way, choices = ['pgd', 'none', 'momentum', 'random'] 109 | OBJ_TYPE: pixel # attack input object, choices = ['pixel', 'pos', 'struc', 'pos+struc', 'pixel+pos', 'pixel+pos+struc'] 110 | LOSS_FUNC: perm # attack loss objective, choices = ['perm', 'ourloss', 'hamming', 'cw', ...] 111 | 112 | # Attacks params for '2step' mode 113 | ATTACK2: 114 | STEP: 1 115 | TYPE: pgd 116 | OBJ_TYPE: pos+struc -------------------------------------------------------------------------------- /experiments/config2.yaml: -------------------------------------------------------------------------------- 1 | MODEL_NAME: vgg16_ngmv2 2 | DATASET_NAME: voc 3 | 4 | DATASET_FULL_NAME: PascalVOC 5 | 6 | MODULE: models.NGM.model_v2 7 | 8 | BACKBONE: VGG16_bn_final 9 | 10 | BATCH_SIZE: 8 11 | DATALOADER_NUM: 2 12 | FP16: False 13 | 14 | RANDOM_SEED: 123 15 | 16 | # available GPU ids 17 | GPUS: 18 | - 0 19 | 20 | # Problem configuration 21 | PROBLEM: 22 | TYPE: 2GM 23 | RESCALE: # rescaled image size 24 | - 256 25 | - 256 26 | 27 | # Graph construction settings 28 | GRAPH: 29 | SRC_GRAPH_CONSTRUCT: tri 30 | TGT_GRAPH_CONSTRUCT: tri 31 | SYM_ADJACENCY: True 32 | 33 | # Training settings 34 | TRAIN: 35 | # start, end epochs 36 | START_EPOCH: 0 37 | NUM_EPOCHS: 16 38 | 39 | MODE: 2step # choices = ['at', '2step'] 40 | LOSS_FUNC: ourloss # choices = ['perm', 'ourloss', 'hamming', ...] 41 | 42 | # for '2step' mode 43 | BURN_IN_PERIOD: 5 44 | 45 | # for AAR regularizer 46 | REG_LEVEL: 1 47 | REG_RATIO: 1.5 48 | 49 | # whether to synchronize optimization way for min-max AT framework, default: false 50 | SYNC_MINMAX: 0 51 | 52 | # set BN mode as evaluation during generating adversarial examples 53 | BN_EVAL: True 54 | 55 | OPTIMIZER: Adam 56 | # learning rate 57 | LR: 2.e-3 58 | SEPARATE_BACKBONE_LR: True 59 | BACKBONE_LR: 2.e-5 60 | MOMENTUM: 0.9 61 | LR_DECAY: 0.5 62 | LR_STEP: # (in epochs) 63 | - 2 64 | - 4 65 | - 6 66 | - 8 67 | - 10 68 | 69 | EPOCH_ITERS: 2000 # iterations per epoch 70 | 71 | CLASS: none 72 | 73 | # Evaluation settings 74 | EVAL: 75 | EPOCH: 15 # where to continue evlauation based on training epochs 76 | SAMPLES: 1000 # number of tested pairs for each class 77 | NUM_EPOCH: 4 # evaluation period during training 78 | 79 | # model parameters 80 | NGM: 81 | EDGE_FEATURE: cat 82 | FEATURE_CHANNEL: 512 83 | SK_ITER_NUM: 20 84 | SK_EPSILON: 1.0e-10 85 | SK_TAU: 0.05 86 | GNN_FEAT: 87 | - 16 88 | - 16 89 | - 16 90 | GNN_LAYER: 3 91 | GAUSSIAN_SIGMA: 1. 92 | SK_EMB: 1 93 | FIRST_ORDER: True 94 | EDGE_EMB: False 95 | 96 | # Attack params for AT 97 | ATTACK: 98 | # basic params 99 | EPSILON_FEATURE: 8 100 | EPSILON_LOCALITY: 8 101 | ALPHA: 1. 102 | STEP: 1 103 | EVAL_ALPHA: 0.25 104 | EVAL_STEP: 10 105 | RESTARTS: 1 106 | EARLY_STOP_RATIO: 0. 107 | # attack attributes 108 | TYPE: pgd # attack optimization way, choices = ['pgd', 'none', 'momentum', 'random'] 109 | OBJ_TYPE: pixel+pos+struc # attack input object, choices = ['pixel', 'pos', 'struc', 'pos+struc', 'pixel+pos', 'pixel+pos+struc'] 110 | LOSS_FUNC: perm # attack loss objective, choices = ['perm', 'ourloss', 'hamming', 'cw', ...] 111 | 112 | # Attacks params for '2step' mode 113 | ATTACK2: 114 | STEP: 1 115 | TYPE: pgd 116 | OBJ_TYPE: pos+struc -------------------------------------------------------------------------------- /experiments/config3.yaml: -------------------------------------------------------------------------------- 1 | MODEL_NAME: vgg16_ngmv2 2 | DATASET_NAME: voc 3 | 4 | DATASET_FULL_NAME: PascalVOC 5 | 6 | MODULE: models.NGM.model_v2 7 | 8 | BACKBONE: VGG16_bn_final 9 | 10 | BATCH_SIZE: 8 11 | DATALOADER_NUM: 2 12 | FP16: False 13 | 14 | RANDOM_SEED: 123 15 | 16 | # available GPU ids 17 | GPUS: 18 | - 0 19 | 20 | # Problem configuration 21 | PROBLEM: 22 | TYPE: 2GM 23 | RESCALE: # rescaled image size 24 | - 256 25 | - 256 26 | 27 | # Graph construction settings 28 | GRAPH: 29 | SRC_GRAPH_CONSTRUCT: tri 30 | TGT_GRAPH_CONSTRUCT: tri 31 | SYM_ADJACENCY: True 32 | 33 | # Training settings 34 | TRAIN: 35 | # start, end epochs 36 | START_EPOCH: 0 37 | NUM_EPOCHS: 16 38 | 39 | MODE: 2step # choices = ['at', '2step'] 40 | LOSS_FUNC: ourloss # choices = ['perm', 'ourloss', 'hamming', ...] 41 | 42 | # for '2step' mode 43 | BURN_IN_PERIOD: 5 44 | 45 | # for AAR regularizer 46 | REG_LEVEL: 1 47 | REG_RATIO: 1.5 48 | 49 | # whether to synchronize optimization way for min-max AT framework, default: false 50 | SYNC_MINMAX: 0 51 | 52 | # set BN mode as evaluation during generating adversarial examples 53 | BN_EVAL: True 54 | 55 | OPTIMIZER: Adam 56 | # learning rate 57 | LR: 2.e-3 58 | SEPARATE_BACKBONE_LR: True 59 | BACKBONE_LR: 2.e-5 60 | MOMENTUM: 0.9 61 | LR_DECAY: 0.5 62 | LR_STEP: # (in epochs) 63 | - 2 64 | - 4 65 | - 6 66 | - 8 67 | - 10 68 | 69 | EPOCH_ITERS: 2000 # iterations per epoch 70 | 71 | CLASS: none 72 | 73 | # Evaluation settings 74 | EVAL: 75 | EPOCH: 15 # where to continue evlauation based on training epochs 76 | SAMPLES: 1000 # number of tested pairs for each class 77 | NUM_EPOCH: 4 # evaluation period during training 78 | 79 | # model parameters 80 | NGM: 81 | EDGE_FEATURE: cat 82 | FEATURE_CHANNEL: 512 83 | SK_ITER_NUM: 20 84 | SK_EPSILON: 1.0e-10 85 | SK_TAU: 0.05 86 | GNN_FEAT: 87 | - 16 88 | - 16 89 | - 16 90 | GNN_LAYER: 3 91 | GAUSSIAN_SIGMA: 1. 92 | SK_EMB: 1 93 | FIRST_ORDER: True 94 | EDGE_EMB: False 95 | 96 | # Attack params for AT 97 | ATTACK: 98 | # basic params 99 | EPSILON_FEATURE: 8 100 | EPSILON_LOCALITY: 8 101 | ALPHA: 1. 102 | STEP: 2 103 | EVAL_ALPHA: 0.25 104 | EVAL_STEP: 10 105 | RESTARTS: 1 106 | EARLY_STOP_RATIO: 0. 107 | # attack attributes 108 | TYPE: pgd # attack optimization way, choices = ['pgd', 'none', 'momentum', 'random'] 109 | OBJ_TYPE: pixel+pos+struc # attack input object, choices = ['pixel', 'pos', 'struc', 'pos+struc', 'pixel+pos', 'pixel+pos+struc'] 110 | LOSS_FUNC: perm # attack loss objective, choices = ['perm', 'ourloss', 'hamming', 'cw', ...] 111 | 112 | # Attacks params for '2step' mode 113 | ATTACK2: 114 | STEP: 1 115 | TYPE: pgd 116 | OBJ_TYPE: pos+struc -------------------------------------------------------------------------------- /src/extension/sparse_dot/csr_dot_csc_cuda.cu: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | #include 4 | #include 5 | 6 | 7 | template 8 | __global__ void csr_dot_csc_cuda_kernel( 9 | const int64_t* __restrict__ t1_indices, 10 | const int64_t* __restrict__ t1_indptr, 11 | const scalar_t* __restrict__ t1_data, 12 | const int64_t* __restrict__ t2_indices, 13 | const int64_t* __restrict__ t2_indptr, 14 | const scalar_t* __restrict__ t2_data, 15 | scalar_t* __restrict__ out_dense, 16 | const int64_t out_h, 17 | const int64_t out_w 18 | ) 19 | { 20 | const int64_t ij = blockIdx.x * blockDim.x + threadIdx.x; 21 | const int64_t b = blockIdx.y; 22 | 23 | if (ij < out_h * out_w) 24 | { 25 | const int64_t i = ij / out_w; 26 | const int64_t j = ij % out_w; 27 | 28 | const int64_t t1_start = t1_indptr[b * out_h + i]; 29 | const int64_t t1_stop = t1_indptr[b * out_h + i + 1]; 30 | 31 | const int64_t t2_start = t2_indptr[b * out_w + j]; 32 | const int64_t t2_stop = t2_indptr[b * out_w + j + 1]; 33 | 34 | scalar_t outp = 0; 35 | int64_t t1_ptr_idx = t1_start; 36 | int64_t t2_ptr_idx = t2_start; 37 | 38 | while (t1_ptr_idx < t1_stop && t2_ptr_idx < t2_stop) 39 | { 40 | int64_t t1_cur_indice = t1_indices[t1_ptr_idx]; 41 | int64_t t2_cur_indice = t2_indices[t2_ptr_idx]; 42 | if (t1_cur_indice == t2_cur_indice) 43 | { 44 | outp += t1_data[t1_ptr_idx] * t2_data[t2_ptr_idx]; 45 | t1_ptr_idx++; 46 | t2_ptr_idx++; 47 | } 48 | else if (t1_cur_indice < t2_cur_indice) 49 | t1_ptr_idx++; 50 | else 51 | t2_ptr_idx++; 52 | } 53 | out_dense[b * out_w * out_h + i * out_w + j] = outp; 54 | } 55 | } 56 | 57 | 58 | at::Tensor csr_dot_csc_cuda( 59 | at::Tensor t1_indices, 60 | at::Tensor t1_indptr, 61 | at::Tensor t1_data, 62 | at::Tensor t2_indices, 63 | at::Tensor t2_indptr, 64 | at::Tensor t2_data, 65 | int64_t batch_size, 66 | int64_t out_h, 67 | int64_t out_w 68 | ){ 69 | auto out_dense = at::zeros({batch_size, out_h, out_w}, t1_data.type()); 70 | 71 | const int threads = 1024; 72 | const dim3 blocks((out_h * out_w + threads - 1) / threads, batch_size); 73 | 74 | AT_DISPATCH_FLOATING_TYPES_AND_HALF(t1_data.type(), "csr_dot_csc_cuda", ([&] { 75 | csr_dot_csc_cuda_kernel<<>>( 76 | t1_indices.data(), 77 | t1_indptr.data(), 78 | t1_data.data(), 79 | t2_indices.data(), 80 | t2_indptr.data(), 81 | t2_data.data(), 82 | out_dense.data(), 83 | out_h, 84 | out_w); 85 | })); 86 | return out_dense; 87 | } 88 | -------------------------------------------------------------------------------- /models/CIE/README.md: -------------------------------------------------------------------------------- 1 | # CIE-H 2 | 3 | Our implementation of the following paper: 4 | * Tianshu Yu, Runzhong Wang, Junchi Yan, Baoxin Li. "Learning deep graph matching with channel-independent embedding and Hungarian attention." _ICLR 2020_. 5 | [[paper]](https://openreview.net/forum?id=rJgBd2NYPH) 6 | 7 | CIE-H follows the CNN-GNN-metric-Sinkhorn pipeline proposed by PCA-GM, and it improves PCA-GM from two aspects: 8 | 1) A channel-independent edge embedding module for better graph feature extraction; 9 | 2) A Hungarian Attention module that dynamically constructs a structured and sparsely connected layer, 10 | taking into account the most contributing matching pairs as hard attention during training. 11 | 12 | ## Benchmark Results 13 | ### PascalVOC - 2GM 14 | 15 | experiment config: ``experiments/vgg16_cie_voc.yaml`` 16 | 17 | pretrained model: [google drive](https://drive.google.com/file/d/1oRwcnw06t1rCbrIN_7p8TJZY-XkBOFEp/view?usp=sharing) 18 | 19 | | model | year | aero | bike | bird | boat | bottle | bus | car | cat | chair | cow | table | dog | horse | mbkie | person | plant | sheep | sofa | train | tv | mean | 20 | | --------------------- | ---- | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | 21 | | [CIE-H](https://thinkmatch.readthedocs.io/en/latest/guide/models.html#cie-h) | 2020 | 0.5250 | 0.6858 | 0.7015 | 0.5706 | 0.8207 | 0.7700 | 0.7073 | 0.7313 | 0.4383 | 0.6994 | 0.6237 | 0.7018 | 0.7031 | 0.6641 | 0.4763 | 0.8525 | 0.7172 | 0.6400 | 0.8385 | 0.9168 | 0.6892 | 22 | 23 | ### Willow Object Class - 2GM 24 | 25 | experiment config: ``experiments/vgg16_cie_willow.yaml`` 26 | 27 | pretrained model: [google drive](https://drive.google.com/file/d/1aUdNTWlFxk-sf-bj08ADUoo9CSIQjzDb/view?usp=sharing) 28 | 29 | | model | year | remark | Car | Duck | Face | Motorbike | Winebottle | mean | 30 | | ------------------------ | ---- | --------------- | ------ | ------ | ------ | --------- | ---------- | ------ | 31 | | [CIE-H](https://thinkmatch.readthedocs.io/en/latest/guide/models.html#cie-h) | 2020 | - | 0.8581 | 0.8206 | 0.9994 | 0.8836 | 0.8871 | 0.8898 | 32 | 33 | 34 | ## File Organization 35 | ``` 36 | ├── model.py 37 | | the implementation of training/evaluation procedures of BBGM 38 | └── model_config.py 39 | the declaration of model hyperparameters 40 | ``` 41 | some files are borrowed from ``models/PCA`` 42 | 43 | ## Credits and Citation 44 | 45 | Please cite the following paper if you use this model in your research: 46 | ``` 47 | @inproceedings{YuICLR20, 48 | title={Learning deep graph matching with channel-independent embedding and Hungarian attention}, 49 | author={Yu, Tianshu and Wang, Runzhong and Yan, Junchi and Li, Baoxin}, 50 | booktitle={International Conference on Learning Representations}, 51 | year={2020} 52 | } 53 | ``` 54 | -------------------------------------------------------------------------------- /models/GMN/affinity_layer.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from torch.nn.parameter import Parameter 4 | from torch import Tensor 5 | import math 6 | 7 | 8 | class InnerpAffinity(nn.Module): 9 | """ 10 | Affinity Layer to compute the affinity matrix via inner product from feature space. 11 | Me = X * Lambda * Y^T 12 | Mp = Ux * Uy^T 13 | Parameter: scale of weight d 14 | Input: edgewise (pairwise) feature X, Y 15 | pointwise (unary) feature Ux, Uy 16 | Output: edgewise affinity matrix Me 17 | pointwise affinity matrix Mp 18 | Weight: weight matrix Lambda = [[Lambda1, Lambda2], 19 | [Lambda2, Lambda1]] 20 | where Lambda1, Lambda2 > 0 21 | """ 22 | def __init__(self, d): 23 | super(InnerpAffinity, self).__init__() 24 | self.d = d 25 | self.lambda1 = Parameter(Tensor(self.d, self.d)) 26 | self.lambda2 = Parameter(Tensor(self.d, self.d)) 27 | self.relu = nn.ReLU() # problem: if weight<0, then always grad=0. So this parameter is never updated! 28 | self.reset_parameters() 29 | 30 | def reset_parameters(self): 31 | stdv = 1. / math.sqrt(self.lambda1.size(1) * 2) 32 | self.lambda1.data.uniform_(-stdv, stdv) 33 | self.lambda2.data.uniform_(-stdv, stdv) 34 | self.lambda1.data += torch.eye(self.d) / 2 35 | self.lambda2.data += torch.eye(self.d) / 2 36 | 37 | def forward(self, X, Y, Ux, Uy, w1=1, w2=1): 38 | assert X.shape[1] == Y.shape[1] == 2 * self.d 39 | lambda1 = self.relu(self.lambda1 + self.lambda1.transpose(0, 1)) * w1 40 | lambda2 = self.relu(self.lambda2 + self.lambda2.transpose(0, 1)) * w2 41 | weight = torch.cat((torch.cat((lambda1, lambda2)), 42 | torch.cat((lambda2, lambda1))), 1) 43 | Me = torch.matmul(X.transpose(1, 2), weight) 44 | Me = torch.matmul(Me, Y) 45 | Mp = torch.matmul(Ux.transpose(1, 2), Uy) 46 | 47 | return Me, Mp 48 | 49 | 50 | class GaussianAffinity(nn.Module): 51 | """ 52 | Affinity Layer to compute the affinity matrix via gaussian kernel from feature space. 53 | Me = exp(- L2(X, Y) / sigma) 54 | Mp = Ux * Uy^T 55 | Parameter: scale of weight d, gaussian kernel sigma 56 | Input: edgewise (pairwise) feature X, Y 57 | pointwise (unary) feature Ux, Uy 58 | Output: edgewise affinity matrix Me 59 | pointwise affinity matrix Mp 60 | """ 61 | 62 | def __init__(self, d, sigma): 63 | super(GaussianAffinity, self).__init__() 64 | self.d = d 65 | self.sigma = sigma 66 | 67 | def forward(self, X, Y, Ux=None, Uy=None, ae=1., ap=1.): 68 | assert X.shape[1] == Y.shape[1] == self.d 69 | 70 | X = X.unsqueeze(-1).expand(*X.shape, Y.shape[2]) 71 | Y = Y.unsqueeze(-2).expand(*Y.shape[:2], X.shape[2], Y.shape[2]) 72 | # dist = torch.sum(torch.pow(torch.mul(X - Y, self.w.unsqueeze(0).unsqueeze(-1).unsqueeze(-1)), 2), dim=1) 73 | dist = torch.sum(torch.pow(X - Y, 2), dim=1) 74 | dist[torch.isnan(dist)] = float("Inf") 75 | Me = torch.exp(- dist / self.sigma) * ae 76 | 77 | if Ux is None or Uy is None: 78 | return Me 79 | else: 80 | Mp = torch.matmul(Ux.transpose(1, 2), Uy) * ap 81 | return Me, Mp 82 | -------------------------------------------------------------------------------- /models/GANN/README.md: -------------------------------------------------------------------------------- 1 | # GANN 2 | 3 | Our implementation of the following paper: 4 | * Runzhong Wang, Junchi Yan and Xiaokang Yang. "Graduated Assignment for Joint Multi-Graph Matching and Clustering with Application to Unsupervised Graph Matching Network Learning." _NeurIPS 2020_. 5 | [[paper]](https://papers.nips.cc/paper/2020/hash/e6384711491713d29bc63fc5eeb5ba4f-Abstract.html) 6 | * Runzhong Wang, Shaofei Jiang, Junchi Yan and Xiaokang Yang. "Robust Self-supervised Learning of Deep Graph Matching with Mixture of Modes." _Submitted to TPAMI_. 7 | [[project page]](https://thinklab.sjtu.edu.cn/project/GANN-GM/index.html) 8 | 9 | GANN proposes a self-supervised learning framework by leveraging graph matching solvers to provide pseudo labels to train the neural network module in deep graph matching pipeline. We propose a general graph matching solver for various graph matching settings based on the classic Graduated Assignment (GA) algorithm. 10 | 11 | The variants on three different graph matching settings are denoted by different suffixes: 12 | * **GANN-2GM**: self-supervised learning graduated assignment neural network for **two-grpah matching** 13 | * **GANN-MGM**: self-supervised learning graduated assignment neural network for **multi-grpah matching** 14 | * **GANN-MGM3**: self-supervised learning graduated assignment neural network for **multi-graph matching with a mixture of modes** (this setting is also known as multi-graph matching and clustering in the NeurIPS paper) 15 | 16 | GANN-MGM notably surpass supervised learning methods on the relatively small dataset Willow Object Class. 17 | 18 | ## Benchmark Results 19 | ### Willow Object Class - MGM 20 | 21 | experiment config: ``experiments/vgg16_gann-mgm_willow.yaml`` 22 | 23 | pretrained model: [google drive](https://drive.google.com/file/d/15Sg6mi9nrpsD4MAjp8b138-t-17VbYsw/view?usp=sharing) 24 | 25 | | model | year | remark | Car | Duck | Face | Motorbike | Winebottle | mean | 26 | | ------------------------ | ---- | --------------- | ------ | ------ | ------ | --------- | ---------- | ------ | 27 | | [GANN-MGM](https://thinkmatch.readthedocs.io/en/latest/guide/models.html#gann) | 2020 | self-supervised | 0.9600 | 0.9642 | 1.0000 | 1.0000 | 0.9879 | 0.9906 | 28 | 29 | ## File Organization 30 | ``` 31 | ├── graduated_assignment.py 32 | | the implementation of the graduated assignment algorithm covering all scenarios 33 | ├── model.py 34 | | the implementation of training/evaluation procedures of GANN-GM/MGM/MGM3 35 | └── model_config.py 36 | the declaration of model hyperparameters 37 | ``` 38 | 39 | ## Credits and Citation 40 | 41 | Please cite the following papers if you use this model in your research: 42 | ``` 43 | @inproceedings{WangNeurIPS20, 44 | author = {Runzhong Wang and Junchi Yan and Xiaokang Yang}, 45 | title = {Graduated Assignment for Joint Multi-Graph Matching and Clustering with Application to Unsupervised Graph Matching Network Learning}, 46 | booktitle = {Neural Information Processing Systems}, 47 | year = {2020} 48 | } 49 | 50 | @unpublished{WangPAMIsub21, 51 | title={Robust Self-supervised Learning of Deep Graph Matching with Mixture of Modes}, 52 | author={Wang, Runzhong and Jiang, Shaofei and Yan, Junchi and Yang, Xiaokang}, 53 | note={submitted to IEEE Transactions of Pattern Analysis and Machine Intelligence}, 54 | year={2021} 55 | } 56 | ``` 57 | -------------------------------------------------------------------------------- /models/GMN/README.md: -------------------------------------------------------------------------------- 1 | # GMN 2 | 3 | Our implementation of the following paper: 4 | * Andrei Zanfir and Cristian Sminchisescu. "Deep Learning of Graph Matching." _CVPR 2018_. 5 | [[paper]](http://openaccess.thecvf.com/content_cvpr_2018/html/Zanfir_Deep_Learning_of_CVPR_2018_paper.html) 6 | 7 | GMN proposes the first deep graph matching pipeline which is end-to-end trainable via supervised learning and gradient descent. It proposes to combine the following components to formulate the graph matching pipeline: 8 | * VGG16 CNN to extract image features 9 | * Delaunay triangulation to build graphs 10 | * Building affinity matrix efficiently via Factorized Graph Matching (FGM) 11 | * Solving the resulting Quadratic Assignment Problem (QAP) by Spectral Matching (SM) and Sinkhorn algorithm which are differentiable 12 | * Supervised learning based on pixel offset loss (known as "Robust loss" in the paper) 13 | 14 | ## Benchmark Results 15 | ### PascalVOC - 2GM 16 | experiment config: ``experiments/vgg16_gmn_voc.yaml`` 17 | 18 | pretrained model: [google drive](https://drive.google.com/file/d/1X8p4XjzqGDniYirwSNqsQhBLWB5VcqN2/view?usp=sharing) 19 | 20 | | model | year | aero | bike | bird | boat | bottle | bus | car | cat | chair | cow | table | dog | horse | mbkie | person | plant | sheep | sofa | train | tv | mean | 21 | | ---------------------- | ---- | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | 22 | | [GMN](https://thinkmatch.readthedocs.io/en/latest/guide/models.html#gmn) | 2018 | 0.4163 | 0.5964 | 0.6027 | 0.4795 | 0.7918 | 0.7020 | 0.6735 | 0.6488 | 0.3924 | 0.6128 | 0.6693 | 0.5976 | 0.6106 | 0.5975 | 0.3721 | 0.7818 | 0.6800 | 0.4993 | 0.8421 | 0.9141 | 0.6240 | 23 | 24 | ### Willow Object Class - 2GM 25 | experiment config: ``experiments/vgg16_gmn_willow.yaml`` 26 | 27 | pretrained model: [google drive](https://drive.google.com/file/d/1PWM1i0oywH3hrwPdYerPazRmhApC0B4U/view?usp=sharing) 28 | 29 | | model | year | remark | Car | Duck | Face | Motorbike | Winebottle | mean | 30 | | ------------------------ | ---- | --------------- | ------ | ------ | ------ | --------- | ---------- | ------ | 31 | | [GMN](https://thinkmatch.readthedocs.io/en/latest/guide/models.html#gmn) | 2018 | - | 0.6790 | 0.7670 | 0.9980 | 0.6920 | 0.8310 | 0.7934 | 32 | 33 | ## File Organization 34 | ``` 35 | ├── affinity_layer.py 36 | | the implementation of affinity layer to compute the affinity matrix for GMN 37 | ├── displacement_layer.py 38 | | the implementation of the displacement layer to compute the pixel offset loss 39 | ├── model.py 40 | | the implementation of training/evaluation procedures of GMN 41 | ├── model_config.py 42 | | the declaration of model hyperparameters 43 | └── voting_layer.py 44 | the implementation of voting layer to compute the row-stotatic matrix with softmax 45 | ``` 46 | 47 | ## Credits and Citation 48 | 49 | Please cite the following paper if you use this model in your research: 50 | ``` 51 | @inproceedings{ZanfirCVPR18, 52 | author = {A. Zanfir and C. Sminchisescu}, 53 | title = {Deep Learning of Graph Matching}, 54 | booktitle = {IEEE Conference on Computer Vision and Pattern Recognition}, 55 | pages={2684--2693}, 56 | year={2018} 57 | } 58 | ``` 59 | -------------------------------------------------------------------------------- /src/utils/data_to_cuda.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from src.sparse_torch.csx_matrix import CSRMatrix3d, CSCMatrix3d 3 | import torch_geometric as pyg 4 | from copy import deepcopy 5 | 6 | def data_to_cuda(inputs): 7 | """ 8 | Call cuda() on all tensor elements in inputs 9 | :param inputs: input list/dictionary 10 | :return: identical to inputs while all its elements are on cuda 11 | """ 12 | if type(inputs) is list: 13 | for i, x in enumerate(inputs): 14 | inputs[i] = data_to_cuda(x) 15 | elif type(inputs) is tuple: 16 | inputs = list(inputs) 17 | for i, x in enumerate(inputs): 18 | inputs[i] = data_to_cuda(x) 19 | elif type(inputs) is dict: 20 | for key in inputs: 21 | inputs[key] = data_to_cuda(inputs[key]) 22 | elif type(inputs) in [str, int, float]: 23 | inputs = inputs 24 | elif type(inputs) in [torch.Tensor, CSRMatrix3d, CSCMatrix3d]: 25 | inputs = inputs.cuda() 26 | elif type(inputs) in [pyg.data.Data, pyg.data.Batch]: 27 | inputs = inputs.to('cuda') 28 | else: 29 | raise TypeError('Unknown type of inputs: {}'.format(type(inputs))) 30 | return inputs 31 | 32 | def cuda_copy(inputs, inputs_att): 33 | """ 34 | Call cuda() on all tensor elements in inputs 35 | :param inputs: input list/dictionary 36 | :return: identical to inputs while all its elements are on cuda 37 | """ 38 | if type(inputs) is list: 39 | for i, x in enumerate(inputs): 40 | inputs[i] = data_to_cuda(x) 41 | elif type(inputs) is tuple: 42 | inputs = list(inputs) 43 | for i, x in enumerate(inputs): 44 | inputs[i] = data_to_cuda(x) 45 | elif type(inputs) is dict: 46 | for key in inputs: 47 | # inputs_att[key] = inputs[key].detach().clone() 48 | inputs_att[key] = data_to_cuda(inputs[key]) 49 | elif type(inputs) in [str, int, float]: 50 | inputs = deepcopy(inputs) 51 | elif type(inputs) in [torch.Tensor, CSRMatrix3d, CSCMatrix3d]: 52 | inputs = inputs.detach().clone() 53 | elif type(inputs) in [pyg.data.Data, pyg.data.Batch]: 54 | inputs = inputs.detach().clone() 55 | else: 56 | raise TypeError('Unknown type of inputs: {}'.format(type(inputs))) 57 | return inputs_att 58 | 59 | def data_to_cuda_sample(inputs, idx_to_fool): 60 | """ 61 | Call cuda() on all tensor elements in inputs 62 | :param inputs: input list/dictionary 63 | :return: identical to inputs while all its elements are on cuda 64 | """ 65 | if type(inputs) is list: 66 | for i, x in enumerate(inputs): 67 | inputs[i] = data_to_cuda_sample(x, idx_to_fool) 68 | elif type(inputs) is tuple: 69 | inputs = list(inputs) 70 | for i, x in enumerate(inputs): 71 | inputs[i] = data_to_cuda_sample(x, idx_to_fool) 72 | elif type(inputs) is dict: 73 | for key in inputs: 74 | # print(key) 75 | inputs[key] = data_to_cuda_sample(inputs[key], idx_to_fool) 76 | elif type(inputs) in [str, int, float]: 77 | inputs = inputs 78 | elif type(inputs) in [torch.Tensor, CSRMatrix3d, CSCMatrix3d]: 79 | # import pdb; pdb.set_trace() 80 | if type(inputs) == torch.Tensor: 81 | inputs = inputs[idx_to_fool].cuda() 82 | else: 83 | inputs = inputs.cuda() 84 | elif type(inputs) in [pyg.data.Data, pyg.data.Batch]: 85 | # import pdb; pdb.set_trace() 86 | inputs = inputs.to('cuda') 87 | else: 88 | raise TypeError('Unknown type of inputs: {}'.format(type(inputs))) 89 | return inputs -------------------------------------------------------------------------------- /src/parallel/scatter_gather.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn.parallel.scatter_gather as torch_ 3 | from src.sparse_torch import CSRMatrix3d, CSCMatrix3d, concatenate 4 | 5 | 6 | def scatter(inputs, target_gpus, dim=0): 7 | """ 8 | Slices tensors into approximately equal chunks and 9 | distributes them across given GPUs. Duplicates 10 | references to objects that are not tensors. 11 | """ 12 | def scatter_map(obj): 13 | if isinstance(obj, torch.Tensor): 14 | return torch_.Scatter.apply(target_gpus, None, dim, obj) 15 | if isinstance(obj, tuple) and len(obj) > 0: 16 | return list(zip(*map(scatter_map, obj))) 17 | if isinstance(obj, list) and len(obj) > 0: 18 | return list(map(list, zip(*map(scatter_map, obj)))) 19 | if isinstance(obj, dict) and len(obj) > 0: 20 | return list(map(type(obj), zip(*map(scatter_map, obj.items())))) 21 | 22 | # modified here 23 | if isinstance(obj, CSRMatrix3d) or isinstance(obj, CSCMatrix3d): 24 | return scatter_sparse_matrix(target_gpus, obj) 25 | 26 | return [obj for targets in target_gpus] 27 | 28 | # After scatter_map is called, a scatter_map cell will exist. This cell 29 | # has a reference to the actual function scatter_map, which has references 30 | # to a closure that has a reference to the scatter_map cell (because the 31 | # fn is recursive). To avoid this reference cycle, we set the function to 32 | # None, clearing the cell 33 | try: 34 | return scatter_map(inputs) 35 | finally: 36 | scatter_map = None 37 | 38 | 39 | def scatter_kwargs(inputs, kwargs, target_gpus, dim=0): 40 | """Scatter with support for kwargs dictionary""" 41 | inputs = scatter(inputs, target_gpus, dim) if inputs else [] 42 | kwargs = scatter(kwargs, target_gpus, dim) if kwargs else [] 43 | if len(inputs) < len(kwargs): 44 | inputs.extend([() for _ in range(len(kwargs) - len(inputs))]) 45 | elif len(kwargs) < len(inputs): 46 | kwargs.extend([{} for _ in range(len(inputs) - len(kwargs))]) 47 | inputs = tuple(inputs) 48 | kwargs = tuple(kwargs) 49 | return inputs, kwargs 50 | 51 | 52 | def scatter_sparse_matrix(target_gpus, obj): 53 | """Scatter for customized sparse matrix""" 54 | def get_device(i): 55 | return torch.device('cuda:{}'.format(i)) if i != -1 else torch.device('cpu') 56 | step = len(obj) // len(target_gpus) 57 | return tuple([obj[i:i+step].to(get_device(i // step)) for i in range(0, len(obj), step)]) 58 | 59 | 60 | def gather(outputs, target_device, dim=0): 61 | """ 62 | Gathers tensors from different GPUs on a specified device (-1 means the CPU). 63 | """ 64 | def gather_map(outputs): 65 | out = outputs[0] 66 | if isinstance(out, torch.Tensor): 67 | return torch_.Gather.apply(target_device, dim, *outputs) 68 | 69 | # modified here 70 | if isinstance(out, CSRMatrix3d) or isinstance(out, CSCMatrix3d): 71 | return concatenate(*outputs, device=target_device) 72 | 73 | if out is None: 74 | return None 75 | if isinstance(out, dict): 76 | if not all((len(out) == len(d) for d in outputs)): 77 | raise ValueError('All dicts must have the same number of keys') 78 | return type(out)(((k, gather_map([d[k] for d in outputs])) 79 | for k in out)) 80 | return type(out)(map(gather_map, zip(*outputs))) 81 | 82 | # Recursive function calls like this create reference cycles. 83 | # Setting the function to None clears the refcycle. 84 | try: 85 | return gather_map(outputs) 86 | finally: 87 | gather_map = None 88 | -------------------------------------------------------------------------------- /src/dataset/dataset_config.py: -------------------------------------------------------------------------------- 1 | from easydict import EasyDict as edict 2 | 3 | __C = edict() 4 | 5 | dataset_cfg = __C 6 | # Pascal VOC 2011 dataset with keypoint annotations 7 | __C.PascalVOC = edict() 8 | __C.PascalVOC.KPT_ANNO_DIR = 'data/PascalVOC/annotations/' # keypoint annotation 9 | __C.PascalVOC.ROOT_DIR = 'data/PascalVOC/VOC2011/' # original VOC2011 dataset 10 | __C.PascalVOC.SET_SPLIT = 'data/PascalVOC/voc2011_pairs.npz' # set split path 11 | __C.PascalVOC.CLASSES = ['aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair', 'cow', 12 | 'diningtable', 'dog', 'horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train', 13 | 'tvmonitor'] 14 | 15 | # Willow-Object Class dataset 16 | __C.WillowObject = edict() 17 | __C.WillowObject.ROOT_DIR = 'data/WILLOW-ObjectClass' 18 | __C.WillowObject.CLASSES = ['Car', 'Duck', 'Face', 'Motorbike', 'Winebottle'] 19 | __C.WillowObject.KPT_LEN = 10 20 | __C.WillowObject.TRAIN_NUM = 20 21 | __C.WillowObject.SPLIT_OFFSET = 0 22 | __C.WillowObject.TRAIN_SAME_AS_TEST = False 23 | __C.WillowObject.RAND_OUTLIER = 0 24 | 25 | # Synthetic dataset 26 | __C.SYNTHETIC = edict() 27 | __C.SYNTHETIC.DIM = 1024 28 | __C.SYNTHETIC.TRAIN_NUM = 100 # training graphs 29 | __C.SYNTHETIC.TEST_NUM = 100 # testing graphs 30 | __C.SYNTHETIC.MIXED_DATA_NUM = 10 # num of samples in mixed synthetic test 31 | __C.SYNTHETIC.RANDOM_EXP_ID = 0 # id of random experiment 32 | __C.SYNTHETIC.EDGE_DENSITY = 0.3 # edge_num = X * node_num^2 / 4 33 | __C.SYNTHETIC.KPT_NUM = 10 # number of nodes (inliers) 34 | __C.SYNTHETIC.OUT_NUM = 0 # number of outliers 35 | __C.SYNTHETIC.FEAT_GT_UNIFORM = 1. # reference node features in uniform(-X, X) for each dimension 36 | __C.SYNTHETIC.FEAT_NOISE_STD = 0.1 # corresponding node features add a random noise ~ N(0, X^2) 37 | __C.SYNTHETIC.POS_GT_UNIFORM = 256. # reference keypoint position in image: uniform(0, X) 38 | __C.SYNTHETIC.POS_AFFINE_DXY = 50. # corresponding position after affine transform: t_x, t_y ~ uniform(-X, X) 39 | __C.SYNTHETIC.POS_AFFINE_S_LOW = 0.8 # corresponding position after affine transform: s ~ uniform(S_LOW, S_HIGH) 40 | __C.SYNTHETIC.POS_AFFINE_S_HIGH = 1.2 41 | __C.SYNTHETIC.POS_AFFINE_DTHETA = 60. # corresponding position after affine transform: theta ~ uniform(-X, X) 42 | __C.SYNTHETIC.POS_NOISE_STD = 10. # corresponding position add a random noise ~ N(0, X^2) after affine transform 43 | 44 | # QAPLIB dataset 45 | __C.QAPLIB = edict() 46 | __C.QAPLIB.DIR = 'data/qapdata' 47 | __C.QAPLIB.FEED_TYPE = 'affmat' # 'affmat' (affinity matrix) or 'adj' (adjacency matrix) 48 | __C.QAPLIB.ONLINE_REPO = 'http://anjos.mgi.polymtl.ca/qaplib/' 49 | __C.QAPLIB.MAX_TRAIN_SIZE = 200 50 | __C.QAPLIB.MAX_TEST_SIZE = 100 51 | 52 | # CUB2011 dataset 53 | __C.CUB2011 = edict() 54 | __C.CUB2011.ROOT_PATH = 'data/CUB_200_2011' 55 | __C.CUB2011.CLASS_SPLIT = 'ori' # choose from 'ori' (original split), 'sup' (super class) or 'all' (all birds as one class) 56 | 57 | # IMC_PT_SparseGM dataset 58 | __C.IMC_PT_SparseGM = edict() 59 | __C.IMC_PT_SparseGM.CLASSES = {'train': ['brandenburg_gate', 'buckingham_palace', 'colosseum_exterior', 60 | 'grand_place_brussels', 'hagia_sophia_interior', 'notre_dame_front_facade', 61 | 'palace_of_westminster', 'pantheon_exterior', 'prague_old_town_square', 62 | 'taj_mahal', 'temple_nara_japan', 'trevi_fountain', 'westminster_abbey'], 63 | 'test': ['reichstag', 'sacre_coeur', 'st_peters_square']} 64 | __C.IMC_PT_SparseGM.ROOT_DIR_NPZ = 'data/IMC_PT_SparseGM/annotation' 65 | __C.IMC_PT_SparseGM.ROOT_DIR_IMG = 'data/IMC_PT_SparseGM/Image_Matching_Challange_Data' 66 | __C.IMC_PT_SparseGM.TOTAL_KPT_NUM = 50 67 | -------------------------------------------------------------------------------- /models/GMN/model.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | 3 | from models.GMN.affinity_layer import InnerpAffinity as Affinity 4 | from src.qap_solvers.spectral_matching import SpectralMatching 5 | from src.qap_solvers.rrwm import RRWM 6 | from src.lap_solvers.sinkhorn import Sinkhorn 7 | from src.lap_solvers.hungarian import hungarian 8 | from src.build_graphs import reshape_edge_feature 9 | from src.feature_align import feature_align 10 | from src.factorize_graph_matching import construct_aff_mat 11 | 12 | from src.utils.config import cfg 13 | 14 | from src.backbone import * 15 | CNN = eval(cfg.BACKBONE) 16 | 17 | 18 | class Net(CNN): 19 | def __init__(self): 20 | super(Net, self).__init__() 21 | self.affinity_layer = Affinity(cfg.GMN.FEATURE_CHANNEL) 22 | if cfg.GMN.GM_SOLVER == 'SM': 23 | self.gm_solver = SpectralMatching(max_iter=cfg.GMN.PI_ITER_NUM, stop_thresh=cfg.GMN.PI_STOP_THRESH) 24 | elif cfg.GMN.GM_SOLVER == 'RRWM': 25 | self.gm_solver = RRWM() 26 | self.sinkhorn = Sinkhorn(max_iter=cfg.GMN.BS_ITER_NUM, tau=1/cfg.GMN.VOTING_ALPHA, epsilon=cfg.GMN.BS_EPSILON, log_forward=False) 27 | self.l2norm = nn.LocalResponseNorm(cfg.GMN.FEATURE_CHANNEL * 2, alpha=cfg.GMN.FEATURE_CHANNEL * 2, beta=0.5, k=0) 28 | self.rescale = cfg.PROBLEM.RESCALE 29 | 30 | def forward(self, data_dict, **kwargs): 31 | if 'images' in data_dict: 32 | # real image data 33 | src, tgt = data_dict['images'] 34 | P_src, P_tgt = data_dict['Ps'] 35 | ns_src, ns_tgt = data_dict['ns'] 36 | G_src, G_tgt = data_dict['Gs'] 37 | H_src, H_tgt = data_dict['Hs'] 38 | K_G, K_H = data_dict['KGHs'] 39 | 40 | # extract feature 41 | src_node = self.node_layers(src) 42 | src_edge = self.edge_layers(src_node) 43 | tgt_node = self.node_layers(tgt) 44 | tgt_edge = self.edge_layers(tgt_node) 45 | 46 | # feature normalization 47 | src_node = self.l2norm(src_node) 48 | src_edge = self.l2norm(src_edge) 49 | tgt_node = self.l2norm(tgt_node) 50 | tgt_edge = self.l2norm(tgt_edge) 51 | 52 | # arrange features 53 | U_src = feature_align(src_node, P_src, ns_src, self.rescale) 54 | F_src = feature_align(src_edge, P_src, ns_src, self.rescale) 55 | U_tgt = feature_align(tgt_node, P_tgt, ns_tgt, self.rescale) 56 | F_tgt = feature_align(tgt_edge, P_tgt, ns_tgt, self.rescale) 57 | elif 'features' in data_dict: 58 | # synthetic data 59 | src, tgt = data_dict['features'] 60 | P_src, P_tgt = data_dict['Ps'] 61 | ns_src, ns_tgt = data_dict['ns'] 62 | G_src, G_tgt = data_dict['Gs'] 63 | H_src, H_tgt = data_dict['Hs'] 64 | K_G, K_H = data_dict['KGHs'] 65 | 66 | U_src = src[:, :src.shape[1] // 2, :] 67 | F_src = src[:, src.shape[1] // 2:, :] 68 | U_tgt = tgt[:, :tgt.shape[1] // 2, :] 69 | F_tgt = tgt[:, tgt.shape[1] // 2:, :] 70 | else: 71 | raise ValueError('Unknown data type for this model.') 72 | 73 | X = reshape_edge_feature(F_src, G_src, H_src) 74 | Y = reshape_edge_feature(F_tgt, G_tgt, H_tgt) 75 | 76 | # affinity layer 77 | Me, Mp = self.affinity_layer(X, Y, U_src, U_tgt) 78 | 79 | M = construct_aff_mat(Me, Mp, K_G, K_H) 80 | 81 | v = self.gm_solver(M, num_src=P_src.shape[1], ns_src=ns_src, ns_tgt=ns_tgt) 82 | s = v.view(v.shape[0], P_tgt.shape[1], -1).transpose(1, 2) 83 | 84 | s = self.sinkhorn(s, ns_src, ns_tgt) 85 | 86 | data_dict.update({ 87 | 'ds_mat': s, 88 | 'perm_mat': hungarian(s, ns_src, ns_tgt), 89 | 'aff_mat': M 90 | }) 91 | return data_dict 92 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Appearance and Structure Aware Robust Deep Visual Graph Matching: Attack, Defense and Beyond 2 | 3 | Code for [CVPR 2022](https://cvpr2022.thecvf.com/) Paper: "Appearance and Structure Aware Robust Deep Visual Graph Matching: Attack, Defense and Beyond" by Qibing Ren, Qingquan Bao, Runzhong Wang, and Junchi Yan. 4 | 5 | 6 | 7 | 8 | ## News 9 | 10 | 03/27/2022 - Our code is released. 11 | 12 | ## Requisite 13 | 14 | The codes are modified based on [ThinkMatch](https://github.com/Thinklab-SJTU/ThinkMatch) and the basic environment settings also follows it. Here we recommend users to utlize Docker for a quick setup of environments. As for maunal configuration, please refer to [ThinkMatch](https://github.com/Thinklab-SJTU/ThinkMatch) for details. 15 | 16 | ### Docker (RECOMMENDED) from _ThinkMatch_ 17 | 18 | 1. We maintain a prebuilt image at [dockerhub](https://hub.docker.com/r/runzhongwang/thinkmatch): ``runzhongwang/thinkmatch:torch1.6.0-cuda10.1-cudnn7-pyg1.6.3``. It can be used by docker or other container runtimes that support docker images e.g. [singularity](https://sylabs.io/singularity/). 19 | 2. We also provide a ``Dockerfile`` to build your own image (you may need ``docker`` and ``nvidia-docker`` installed on your computer). 20 | 21 | ## What is in this repository 22 | 23 | 1. `train_eval.py` is about the robust training pipeline while `eval.py` is the evaluation codes. 24 | 2. For `attack_utils.py`, it defines the class `AttackGM` that implements our locality attack including several attack baselines. 25 | 3. Moreover, `src/loss_func.py` implements our regularization loss via the parent class `GMLoss`. 26 | 4. `src/utils/config.py` defines a global hyper-parameter dictionary `cfg`, which is referenced everywhere in this project. 27 | 28 | ## Run the Experiment 29 | 30 | Run training and evaluation 31 | 32 | ```bash 33 | python train_eval.py --cfg path/to/your/yaml 34 | ``` 35 | 36 | and replace ``path/to/your/yaml`` by path to your configuration file. For example, to reproduce the ASAR-GM (config 1): 37 | 38 | ```bash 39 | python train_eval.py --cfg experiments/config1.yaml 40 | ``` 41 | 42 | For reproducibility, we release the three configurations of our ASAR-GM in ``experiments/``, namely ``config1.yaml``, ``config2.yaml``, and ``config3.yaml`` respectively. 43 | 44 | To perform various while-box attacks shown in Paper, run the fllowing script: 45 | 46 | ```bash 47 | python train_eval.py --cfg experiments/eval.yaml 48 | ``` 49 | 50 | Note that white-box attack evaluation can be automatically performed by setting 51 | ```yaml 52 | EVAL.MODE: all 53 | ``` 54 | To customize your attack, please change the value of ``EVAL.MODE`` as ``single``. 55 | 56 | Additionally, to perform various black-box attacks shown in Paper, run the fllowing script: 57 | 58 | ```bash 59 | python eval.py --cfg experiments/eval_blackbox.yaml --black 60 | ``` 61 | 62 | Note that you need to specify the model path to the variable ``PRETRAINED_PATH`` for model parameters being loaded. Your are welcome to try your own configurations. If you find a better yaml configuration, please let us know by raising an issue or a PR and we will update the benchmark! 63 | 64 | 65 | ## Pretrained Models 66 | 67 | _RobustMatch_ provides pretrained models of the three configurations of ASAR-GM. The model weights are available via [google drive](https://drive.google.com/drive/folders/1--oZxLn8Oo_JwL_V7QAJ8li1FVuCNqYg?usp=sharing). 68 | 69 | To use the pretrained models, firstly download the weight files, then add the following line to your yaml file: 70 | 71 | ```yaml 72 | PRETRAINED_PATH: path/to/your/pretrained/weights 73 | ``` 74 | 75 | ## Citing this work 76 | 77 | ``` 78 | @inproceedings{ren2022appearance, 79 | title={Appearance and Structure Aware Robust Deep Visual Graph Matching: Attack, Defense and Beyond}, 80 | author={Qibing Ren and Qingquan Bao and Runzhong Wang and Junchi Yan}, 81 | booktitle={CVPR}, 82 | year={2022} 83 | } 84 | ``` 85 | -------------------------------------------------------------------------------- /src/backbone.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from torchvision import models 4 | 5 | 6 | class VGG16_base(nn.Module): 7 | r""" 8 | The base class of VGG16. It downloads the pretrained weight by torchvision API, and maintain the layers needed for 9 | deep graph matching models. 10 | """ 11 | def __init__(self, batch_norm=True, final_layers=False): 12 | super(VGG16_base, self).__init__() 13 | self.node_layers, self.edge_layers, self.final_layers = self.get_backbone(batch_norm) 14 | if not final_layers: self.final_layers = None 15 | self.backbone_params = list(self.parameters()) 16 | 17 | def forward(self, *input): 18 | raise NotImplementedError 19 | 20 | @property 21 | def device(self): 22 | return next(self.parameters()).device 23 | 24 | @staticmethod 25 | def get_backbone(batch_norm): 26 | """ 27 | Get pretrained VGG16 models for feature extraction. 28 | 29 | :return: feature sequence 30 | """ 31 | if batch_norm: 32 | model = models.vgg16_bn(pretrained=True) 33 | else: 34 | model = models.vgg16(pretrained=True) 35 | 36 | conv_layers = nn.Sequential(*list(model.features.children())) 37 | 38 | conv_list = node_list = edge_list = [] 39 | 40 | # get the output of relu4_2(node features) and relu5_1(edge features) 41 | cnt_m, cnt_r = 1, 0 42 | for layer, module in enumerate(conv_layers): 43 | if isinstance(module, nn.Conv2d): 44 | cnt_r += 1 45 | if isinstance(module, nn.MaxPool2d): 46 | cnt_r = 0 47 | cnt_m += 1 48 | conv_list += [module] 49 | 50 | #if cnt_m == 4 and cnt_r == 2 and isinstance(module, nn.ReLU): 51 | if cnt_m == 4 and cnt_r == 3 and isinstance(module, nn.Conv2d): 52 | node_list = conv_list 53 | conv_list = [] 54 | #elif cnt_m == 5 and cnt_r == 1 and isinstance(module, nn.ReLU): 55 | elif cnt_m == 5 and cnt_r == 2 and isinstance(module, nn.Conv2d): 56 | edge_list = conv_list 57 | conv_list = [] 58 | 59 | assert len(node_list) > 0 and len(edge_list) > 0 60 | 61 | # Set the layers as a nn.Sequential module 62 | node_layers = nn.Sequential(*node_list) 63 | edge_layers = nn.Sequential(*edge_list) 64 | final_layers = nn.Sequential(*conv_list, nn.AdaptiveMaxPool2d((1, 1), return_indices=False)) # this final layer follows Rolink et al. ECCV20 65 | 66 | return node_layers, edge_layers, final_layers 67 | 68 | 69 | class VGG16_bn_final(VGG16_base): 70 | r""" 71 | VGG16 with batch normalization and final layers. 72 | """ 73 | def __init__(self): 74 | super(VGG16_bn_final, self).__init__(True, True) 75 | 76 | 77 | class VGG16_bn(VGG16_base): 78 | r""" 79 | VGG16 with batch normalization, without final layers. 80 | """ 81 | def __init__(self): 82 | super(VGG16_bn, self).__init__(True, False) 83 | 84 | 85 | class VGG16_final(VGG16_base): 86 | r""" 87 | VGG16 without batch normalization, with final layers. 88 | """ 89 | def __init__(self): 90 | super(VGG16_final, self).__init__(False, True) 91 | 92 | 93 | class VGG16(VGG16_base): 94 | r""" 95 | VGG16 without batch normalization or final layers. 96 | """ 97 | def __init__(self): 98 | super(VGG16, self).__init__(False, False) 99 | 100 | 101 | class NoBackbone(nn.Module): 102 | r""" 103 | A model with no CNN backbone for non-image data. 104 | """ 105 | def __init__(self, *args, **kwargs): 106 | super(NoBackbone, self).__init__() 107 | self.node_layers, self.edge_layers = None, None 108 | 109 | def forward(self, *input): 110 | raise NotImplementedError 111 | 112 | @property 113 | def device(self): 114 | return next(self.parameters()).device 115 | -------------------------------------------------------------------------------- /models/BBGM/sconv_archs.py: -------------------------------------------------------------------------------- 1 | import torch.nn 2 | import torch 3 | import torch.nn.functional as F 4 | from torch_geometric.nn import SplineConv 5 | 6 | 7 | class SConv(torch.nn.Module): 8 | def __init__(self, input_features, output_features): 9 | super(SConv, self).__init__() 10 | 11 | self.in_channels = input_features 12 | self.num_layers = 2 13 | self.convs = torch.nn.ModuleList() 14 | 15 | for _ in range(self.num_layers): 16 | conv = SplineConv(input_features, output_features, dim=2, kernel_size=5, aggr="max") 17 | self.convs.append(conv) 18 | input_features = output_features 19 | 20 | input_features = output_features 21 | self.out_channels = input_features 22 | self.reset_parameters() 23 | 24 | def reset_parameters(self): 25 | for conv in self.convs: 26 | conv.reset_parameters() 27 | 28 | def forward(self, data): 29 | x, edge_index, edge_attr = data.x, data.edge_index, data.edge_attr 30 | xs = [x] 31 | 32 | for conv in self.convs[:-1]: 33 | xs += [F.relu(conv(xs[-1], edge_index, edge_attr))] 34 | 35 | xs += [self.convs[-1](xs[-1], edge_index, edge_attr)] 36 | return xs[-1] 37 | 38 | 39 | class SiameseSConvOnNodes(torch.nn.Module): 40 | def __init__(self, input_node_dim): 41 | super(SiameseSConvOnNodes, self).__init__() 42 | self.num_node_features = input_node_dim 43 | self.mp_network = SConv(input_features=self.num_node_features, output_features=self.num_node_features) 44 | 45 | def forward(self, graph): 46 | old_features = graph.x 47 | result = self.mp_network(graph) 48 | graph.x = old_features + 0.1 * result 49 | return graph 50 | 51 | 52 | class SiameseNodeFeaturesToEdgeFeatures(torch.nn.Module): 53 | def __init__(self, total_num_nodes): 54 | super(SiameseNodeFeaturesToEdgeFeatures, self).__init__() 55 | self.num_edge_features = total_num_nodes 56 | 57 | def forward(self, graph, hyperedge=False): 58 | orig_graphs = graph.to_data_list() 59 | orig_graphs = [self.vertex_attr_to_edge_attr(graph) for graph in orig_graphs] 60 | if hyperedge: 61 | orig_graphs = [self.vertex_attr_to_hyperedge_attr(graph) for graph in orig_graphs] 62 | return orig_graphs 63 | 64 | def vertex_attr_to_edge_attr(self, graph): 65 | """Assigns the difference of node features to each edge""" 66 | flat_edges = graph.edge_index.transpose(0, 1).reshape(-1) 67 | vertex_attrs = torch.index_select(graph.x, dim=0, index=flat_edges) 68 | 69 | new_shape = (graph.edge_index.shape[1], 2, vertex_attrs.shape[1]) 70 | vertex_attrs_reshaped = vertex_attrs.reshape(new_shape).transpose(0, 1) 71 | new_edge_attrs = vertex_attrs_reshaped[0] - vertex_attrs_reshaped[1] 72 | graph.edge_attr = new_edge_attrs 73 | return graph 74 | 75 | def vertex_attr_to_hyperedge_attr(self, graph): 76 | """Assigns the angle of node features to each hyperedge. 77 | graph.hyperedge_index is the incidence matrix.""" 78 | flat_edges = graph.hyperedge_index.transpose(0, 1).reshape(-1) 79 | vertex_attrs = torch.index_select(graph.x, dim=0, index=flat_edges) 80 | 81 | new_shape = (graph.hyperedge_index.shape[1], 3, vertex_attrs.shape[1]) 82 | 83 | vertex_attrs_reshaped = vertex_attrs.reshape(new_shape).transpose(0, 1) 84 | v01 = vertex_attrs_reshaped[0] - vertex_attrs_reshaped[1] 85 | v02 = vertex_attrs_reshaped[0] - vertex_attrs_reshaped[2] 86 | v12 = vertex_attrs_reshaped[1] - vertex_attrs_reshaped[2] 87 | nv01 = torch.norm(v01, p=2, dim=-1) 88 | nv02 = torch.norm(v02, p=2, dim=-1) 89 | nv12 = torch.norm(v12, p=2, dim=-1) 90 | 91 | cos1 = torch.sum(v01 * v02, dim=-1) / (nv01 * nv02) 92 | cos2 = torch.sum(-v01 * v12, dim=-1) / (nv01 * nv12) 93 | cos3 = torch.sum(-v12 * -v02, dim=-1) / (nv12 * nv02) 94 | 95 | graph.hyperedge_attr = torch.stack((cos1, cos2, cos3), dim=-1) 96 | return graph 97 | -------------------------------------------------------------------------------- /models/BBGM/README.md: -------------------------------------------------------------------------------- 1 | # BBGM 2 | 3 | Our implementation of the following paper: 4 | * Michal Rolínek, Paul Swoboda, Dominik Zietlow, Anselm Paulus, Vít Musil, Georg Martius. "Deep Graph Matching via Blackbox Differentiation of Combinatorial Solvers." _ECCV 2020_. 5 | [[paper]](https://www.ecva.net/papers/eccv_2020/papers_ECCV/papers/123730409.pdf) 6 | 7 | BBGM proposes a new feature extractor by using the global feature of VGG16 and a Spline Convolution module, and such an improved backbone is found effective for image matching problems. 8 | NOTE: Spline convolution is officially named as [SplineCNN](https://arxiv.org/abs/1711.08920), however, since the term "CNN" is conventionally used for image feature extractors, and SplineCNN works on graphs, here we name it as "Spline Convolution" for disambiguation. 9 | 10 | The resulting quadratic assignment problem is solved by a discrete [LPMP solver](https://github.com/LPMP/LPMP), and the gradient is approximated by the [black-box combinatorial solver technique](https://arxiv.org/abs/1912.02175). 11 | 12 | ## Benchmark Results 13 | ### PascalVOC - 2GM 14 | 15 | experiment config: ``experiments/vgg16_bbgm_voc.yaml`` 16 | 17 | pretrained model: [google drive](https://drive.google.com/file/d/1RxC7daviZf3kz2Nvr76DldR_oMHfNB4h/view?usp=sharing) 18 | 19 | | model | year | aero | bike | bird | boat | bottle | bus | car | cat | chair | cow | table | dog | horse | mbkie | person | plant | sheep | sofa | train | tv | mean | 20 | | ---------------------- | ---- | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | 21 | | [BBGM](https://thinkmatch.readthedocs.io/en/latest/guide/models.html#bbgm) | 2020 | 0.6187 | 0.7106 | 0.7969 | 0.7896 | 0.8740 | 0.9401 | 0.8947 | 0.8022 | 0.5676 | 0.7914 | 0.6458 | 0.7892 | 0.7615 | 0.7512 | 0.6519 | 0.9818 | 0.7729 | 0.7701 | 0.9494 | 0.9393 | 0.7899 | 22 | 23 | ### Willow Object Class - 2GM & MGM 24 | 25 | experiment config: ``experiments/vgg16_bbgm_willow.yaml`` 26 | 27 | pretrained model: [google drive](https://drive.google.com/file/d/1bt8wBeimM0ofm3QWEVOWWKxoIVRfFwi-/view?usp=sharing) 28 | 29 | | model | year | remark | Car | Duck | Face | Motorbike | Winebottle | mean | 30 | | ------------------------ | ---- | --------------- | ------ | ------ | ------ | --------- | ---------- | ------ | 31 | | [BBGM](https://thinkmatch.readthedocs.io/en/latest/guide/models.html#bbgm) | 2020 | - | 0.9680 | 0.8990 | 1.0000 | 0.9980 | 0.9940 | 0.9718 | 32 | 33 | ## File Organization 34 | ``` 35 | ├── affinity_layer.py 36 | | the implementation of the inner-product with weight affinity layer proposed by BBGM 37 | ├── model.py 38 | | the implementation of training/evaluation procedures of BBGM 39 | ├── model_config.py 40 | | the declaration of model hyperparameters 41 | └── sconv_archs.py 42 | the implementation of spline convolution (SpilneCNN) operations 43 | ``` 44 | 45 | ## Remarks 46 | It is worth noting that our reproduced result of BBGM is different from the result from their paper. By looking into the code released by BBGM, we find that the authors of BBGM filter out keypoints which are out of the bounding box but we do not. Therefore, the number of nodes of graphs in BBGM is smaller than ours, and the graph matching problem is less challenging than ours. In this repository, we modify BBGM to fit into our setting, and we report our reproduced result for fair comparison. 47 | 48 | ## Credits and Citation 49 | This code is developed based on the [official implementation of BBGM](https://github.com/martius-lab/blackbox-deep-graph-matching). The code is modified to fit our general framework. 50 | 51 | Please cite the following paper if you use this model in your research: 52 | ``` 53 | @inproceedings{RolinekECCV20, 54 | title={Deep Graph Matching via Blackbox Differentiation of Combinatorial Solvers}, 55 | author={Rol{\'\i}nek, Michal and Swoboda, Paul and Zietlow, Dominik and Paulus, Anselm and Musil, V{\'\i}t and Martius, Georg}, 56 | booktitle={European Conference on Computer Vision}, 57 | pages={407--424}, 58 | year={2020}, 59 | organization={Springer} 60 | } 61 | ``` 62 | -------------------------------------------------------------------------------- /models/PCA/README.md: -------------------------------------------------------------------------------- 1 | # PCA-GM 2 | 3 | Our implementation of the following papers: 4 | * Runzhong Wang, Junchi Yan and Xiaokang Yang. "Combinatorial Learning of Robust Deep Graph Matching: an Embedding based Approach." _TPAMI 2020_. 5 | [[paper]](https://ieeexplore.ieee.org/abstract/document/9128045/), [[project page]](https://thinklab.sjtu.edu.cn/IPCA_GM.html) 6 | * Runzhong Wang, Junchi Yan and Xiaokang Yang. "Learning Combinatorial Embedding Networks for Deep Graph Matching." _ICCV 2019_. 7 | [[paper]](http://openaccess.thecvf.com/content_ICCV_2019/papers/Wang_Learning_Combinatorial_Embedding_Networks_for_Deep_Graph_Matching_ICCV_2019_paper.pdf) 8 | 9 | 10 | **PCA-GM** proposes the first deep graph matching network based on graph embedding, and it is composed of the following components: 11 | * VGG16 CNN to extract image features 12 | * Delaunay triangulation to build graphs 13 | * Graph Convolutional Network (GCN) with cross-graph convolution to embed graph structure features 14 | * Solving the resulting Linear Assignment Problem by Sinkhorn network 15 | * Supervised learning based on cross entropy loss (known as "permutation loss" in this paper) 16 | 17 | Such a CNN-GNN-Sinkhorn-CrossEntropy framework has be adopted by many following papers. 18 | 19 | **PCA-GM** is proposed in the conference version. In the journal version, we propose **IPCA-GM**, whereby the cross-graph convolution step is implemented in an iterative manner. 20 | The motivation of the iterative update scheme is that: better embedding features will lead to better cross-graph weight matrix and vice versa. 21 | 22 | ## Benchmark Results 23 | ### PascalVOC - 2GM 24 | * PCA-GM 25 | 26 | experiment config: ``experiments/vgg16_pca_voc.yaml`` 27 | 28 | pretrained model: [google drive](https://drive.google.com/file/d/1JnX3cSPvRYBSrDKVwByzp7CADgVCJCO_/view?usp=sharing) 29 | 30 | * IPCA-GM 31 | 32 | experiment config: ``experiments/vgg16_ipca_voc.yaml`` 33 | 34 | pretrained model: [google drive](https://drive.google.com/file/d/1TGrbSQRmUkClH3Alz2OCwqjl8r8gf5yI/view?usp=sharing) 35 | 36 | | model | year | aero | bike | bird | boat | bottle | bus | car | cat | chair | cow | table | dog | horse | mbkie | person | plant | sheep | sofa | train | tv | mean | 37 | | ---------------------- | ---- | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | ------ | 38 | | [PCA-GM](https://thinkmatch.readthedocs.io/en/latest/guide/models.html#pca-gm) | 2019 | 0.4979 | 0.6193 | 0.6531 | 0.5715 | 0.7882 | 0.7556 | 0.6466 | 0.6969 | 0.4164 | 0.6339 | 0.5073 | 0.6705 | 0.6671 | 0.6164 | 0.4447 | 0.8116 | 0.6782 | 0.5922 | 0.7845 | 0.9042 | 0.6478 | 39 | | [IPCA-GM](https://thinkmatch.readthedocs.io/en/latest/guide/models.html#pca-gm) | 2020 | 0.5378 | 0.6622 | 0.6714 | 0.6120 | 0.8039 | 0.7527 | 0.7255 | 0.7252 | 0.4455 | 0.6524 | 0.5430 | 0.6724 | 0.6790 | 0.6421 | 0.4793 | 0.8435 | 0.7079 | 0.6398 | 0.8380 | 0.9083 | 0.6770 | 40 | 41 | ### Willow Object Class - 2GM 42 | * PCA-GM 43 | 44 | experiment config: ``experiments/vgg16_pca_willow.yaml`` 45 | 46 | pretrained model: [google drive](https://drive.google.com/file/d/1BYFevb7C1mUW9vK-L9wOo0Omtp4V15Ub/view?usp=sharing) 47 | 48 | * IPCA-GM 49 | 50 | experiment config: ``experiments/vgg16_ipca_willow.yaml`` 51 | 52 | pretrained model: [google drive](https://drive.google.com/file/d/1-OcLEwlKiudxs3KoKbFW56kzspqFsoWH/view?usp=sharing) 53 | 54 | | model | year | remark | Car | Duck | Face | Motorbike | Winebottle | mean | 55 | | ------------------------ | ---- | --------------- | ------ | ------ | ------ | --------- | ---------- | ------ | 56 | | [PCA-GM](https://thinkmatch.readthedocs.io/en/latest/guide/models.html#pca-gm) | 2019 | - | 0.8760 | 0.8360 | 1.0000 | 0.7760 | 0.8840 | 0.8744 | 57 | | [IPCA-GM](https://thinkmatch.readthedocs.io/en/latest/guide/models.html#pca) | 2020 | - | 0.9040 | 0.8860 | 1.0000 | 0.8300 | 0.8830 | 0.9006 | 58 | 59 | ## File Organization 60 | ``` 61 | ├── affinity_layer.py 62 | | the implementation of affinity layer to compute the affinity matrix for PCA-GM and IPCA-GM 63 | ├── model.py 64 | | the implementation of training/evaluation procedures of PCA-GM and IPCA-GM 65 | └── model_config.py 66 | the declaration of model hyperparameters 67 | ``` 68 | 69 | ## Credits and Citation 70 | 71 | Please cite the following paper if you use this model in your research: 72 | ``` 73 | @inproceedings{WangICCV19, 74 | author = {Wang, Runzhong and Yan, Junchi and Yang, Xiaokang}, 75 | title = {Learning Combinatorial Embedding Networks for Deep Graph Matching}, 76 | booktitle = {IEEE International Conference on Computer Vision}, 77 | pages={3056--3065}, 78 | year = {2019} 79 | } 80 | ``` 81 | -------------------------------------------------------------------------------- /models/PCA/model.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | from src.lap_solvers.sinkhorn import Sinkhorn 5 | from src.feature_align import feature_align 6 | from src.gconv import Siamese_Gconv 7 | from models.PCA.affinity_layer import Affinity 8 | from src.lap_solvers.hungarian import hungarian 9 | 10 | from src.utils.config import cfg 11 | from models.PCA.model_config import model_cfg 12 | 13 | from src.backbone import * 14 | CNN = eval(cfg.BACKBONE) 15 | 16 | 17 | class Net(CNN): 18 | def __init__(self): 19 | super(Net, self).__init__() 20 | self.sinkhorn = Sinkhorn(max_iter=cfg.PCA.SK_ITER_NUM, epsilon=cfg.PCA.SK_EPSILON, tau=cfg.PCA.SK_TAU) 21 | self.l2norm = nn.LocalResponseNorm(cfg.PCA.FEATURE_CHANNEL * 2, alpha=cfg.PCA.FEATURE_CHANNEL * 2, beta=0.5, k=0) 22 | self.gnn_layer = cfg.PCA.GNN_LAYER 23 | #self.pointer_net = PointerNet(cfg.PCA.GNN_FEAT, cfg.PCA.GNN_FEAT // 2, alpha=cfg.PCA.VOTING_ALPHA) 24 | for i in range(self.gnn_layer): 25 | if i == 0: 26 | gnn_layer = Siamese_Gconv(cfg.PCA.FEATURE_CHANNEL * 2, cfg.PCA.GNN_FEAT) 27 | else: 28 | gnn_layer = Siamese_Gconv(cfg.PCA.GNN_FEAT, cfg.PCA.GNN_FEAT) 29 | self.add_module('gnn_layer_{}'.format(i), gnn_layer) 30 | self.add_module('affinity_{}'.format(i), Affinity(cfg.PCA.GNN_FEAT)) 31 | if i == self.gnn_layer - 2: # only second last layer will have cross-graph module 32 | self.add_module('cross_graph_{}'.format(i), nn.Linear(cfg.PCA.GNN_FEAT * 2, cfg.PCA.GNN_FEAT)) 33 | self.cross_iter = cfg.PCA.CROSS_ITER 34 | self.cross_iter_num = cfg.PCA.CROSS_ITER_NUM 35 | self.rescale = cfg.PROBLEM.RESCALE 36 | 37 | def reload_backbone(self): 38 | self.node_layers, self.edge_layers = self.get_backbone(True) 39 | 40 | 41 | def forward(self, data_dict, **kwargs): 42 | if 'images' in data_dict: 43 | # real image data 44 | src, tgt = data_dict['images'] 45 | P_src, P_tgt = data_dict['Ps'] 46 | ns_src, ns_tgt = data_dict['ns'] 47 | A_src, A_tgt = data_dict['As'] 48 | 49 | # extract feature 50 | src_node = self.node_layers(src) 51 | src_edge = self.edge_layers(src_node) 52 | tgt_node = self.node_layers(tgt) 53 | tgt_edge = self.edge_layers(tgt_node) 54 | 55 | # feature normalization 56 | src_node = self.l2norm(src_node) 57 | src_edge = self.l2norm(src_edge) 58 | tgt_node = self.l2norm(tgt_node) 59 | tgt_edge = self.l2norm(tgt_edge) 60 | 61 | # arrange features 62 | U_src = feature_align(src_node, P_src, ns_src, self.rescale) 63 | F_src = feature_align(src_edge, P_src, ns_src, self.rescale) 64 | U_tgt = feature_align(tgt_node, P_tgt, ns_tgt, self.rescale) 65 | F_tgt = feature_align(tgt_edge, P_tgt, ns_tgt, self.rescale) 66 | elif 'features' in data_dict: 67 | # synthetic data 68 | src, tgt = data_dict['features'] 69 | ns_src, ns_tgt = data_dict['ns'] 70 | A_src, A_tgt = data_dict['As'] 71 | 72 | U_src = src[:, :src.shape[1] // 2, :] 73 | F_src = src[:, src.shape[1] // 2:, :] 74 | U_tgt = tgt[:, :tgt.shape[1] // 2, :] 75 | F_tgt = tgt[:, tgt.shape[1] // 2:, :] 76 | else: 77 | raise ValueError('Unknown data type for this model.') 78 | 79 | emb1, emb2 = torch.cat((U_src, F_src), dim=1).transpose(1, 2), torch.cat((U_tgt, F_tgt), dim=1).transpose(1, 2) 80 | ss = [] 81 | 82 | if not self.cross_iter: 83 | # Vanilla PCA-GM 84 | for i in range(self.gnn_layer): 85 | gnn_layer = getattr(self, 'gnn_layer_{}'.format(i)) 86 | emb1, emb2 = gnn_layer([A_src, emb1], [A_tgt, emb2]) 87 | affinity = getattr(self, 'affinity_{}'.format(i)) 88 | s = affinity(emb1, emb2) 89 | s = self.sinkhorn(s, ns_src, ns_tgt, dummy_row=True) 90 | 91 | ss.append(s) 92 | 93 | if i == self.gnn_layer - 2: 94 | cross_graph = getattr(self, 'cross_graph_{}'.format(i)) 95 | new_emb1 = cross_graph(torch.cat((emb1, torch.bmm(s, emb2)), dim=-1)) 96 | new_emb2 = cross_graph(torch.cat((emb2, torch.bmm(s.transpose(1, 2), emb1)), dim=-1)) 97 | emb1 = new_emb1 98 | emb2 = new_emb2 99 | else: 100 | # IPCA-GM 101 | for i in range(self.gnn_layer - 1): 102 | gnn_layer = getattr(self, 'gnn_layer_{}'.format(i)) 103 | emb1, emb2 = gnn_layer([A_src, emb1], [A_tgt, emb2]) 104 | 105 | emb1_0, emb2_0 = emb1, emb2 106 | s = torch.zeros(emb1.shape[0], emb1.shape[1], emb2.shape[1], device=emb1.device) 107 | 108 | for x in range(self.cross_iter_num): 109 | i = self.gnn_layer - 2 110 | cross_graph = getattr(self, 'cross_graph_{}'.format(i)) 111 | emb1 = cross_graph(torch.cat((emb1_0, torch.bmm(s, emb2_0)), dim=-1)) 112 | emb2 = cross_graph(torch.cat((emb2_0, torch.bmm(s.transpose(1, 2), emb1_0)), dim=-1)) 113 | 114 | i = self.gnn_layer - 1 115 | gnn_layer = getattr(self, 'gnn_layer_{}'.format(i)) 116 | emb1, emb2 = gnn_layer([A_src, emb1], [A_tgt, emb2]) 117 | affinity = getattr(self, 'affinity_{}'.format(i)) 118 | s = affinity(emb1, emb2) 119 | s = self.sinkhorn(s, ns_src, ns_tgt, dummy_row=True) 120 | ss.append(s) 121 | 122 | data_dict.update({ 123 | 'ds_mat': ss[-1], 124 | 'perm_mat': hungarian(ss[-1], ns_src, ns_tgt) 125 | }) 126 | return data_dict 127 | -------------------------------------------------------------------------------- /models/PCA/affinity_layer.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from torch.nn.parameter import Parameter 4 | from torch import Tensor 5 | import math 6 | 7 | 8 | class Affinity(nn.Module): 9 | """ 10 | Affinity Layer to compute the affinity matrix from feature space. 11 | M = X * A * Y^T 12 | Parameter: scale of weight d 13 | Input: feature X, Y 14 | Output: affinity matrix M 15 | """ 16 | def __init__(self, d): 17 | super(Affinity, self).__init__() 18 | self.d = d 19 | self.A = Parameter(Tensor(self.d, self.d)) 20 | self.reset_parameters() 21 | 22 | def reset_parameters(self): 23 | stdv = 1. / math.sqrt(self.d) 24 | self.A.data.uniform_(-stdv, stdv) 25 | self.A.data += torch.eye(self.d) 26 | 27 | def forward(self, X, Y): 28 | assert X.shape[2] == Y.shape[2] == self.d 29 | M = torch.matmul(X, self.A) 30 | #M = torch.matmul(X, (self.A + self.A.transpose(0, 1)) / 2) 31 | M = torch.matmul(M, Y.transpose(1, 2)) 32 | return M 33 | 34 | class AffinityInp(nn.Module): 35 | """ 36 | Affinity Layer to compute inner product affinity matrix from feature space. 37 | M = X * A * Y^T 38 | Parameter: scale of weight d 39 | Input: feature X, Y 40 | Output: affinity matrix M 41 | """ 42 | def __init__(self, d): 43 | super(AffinityInp, self).__init__() 44 | self.d = d 45 | 46 | def forward(self, X, Y): 47 | assert X.shape[2] == Y.shape[2] == self.d 48 | M = torch.matmul(X, Y.transpose(1, 2)) 49 | return M 50 | 51 | 52 | 53 | class AffinityLR(nn.Module): 54 | def __init__(self, d, k=512): 55 | super(AffinityLR, self).__init__() 56 | self.d = d 57 | self.k = k 58 | self.A = Parameter(Tensor(self.d, self.k)) 59 | self.relu = nn.ReLU() 60 | self.reset_parameters() 61 | 62 | def reset_parameters(self): 63 | stdv = 1. / math.sqrt(self.d) 64 | self.A.data.uniform_(-stdv, stdv) 65 | 66 | def forward(self, X, Y): 67 | assert X.shape[2] == Y.shape[2] == self.d 68 | M = torch.matmul(self.A, self.A.transpose(0, 1)) 69 | M = torch.matmul(X, M) 70 | M = torch.matmul(M, Y.transpose(1, 2)) 71 | 72 | return self.relu(M.squeeze()) 73 | 74 | class AffinityMah(nn.Module): 75 | def __init__(self, d, k=100): 76 | super(AffinityMah, self).__init__() 77 | self.d = d 78 | self.k = k 79 | self.A = Parameter(Tensor(self.d, self.k)) 80 | self.relu = nn.ReLU() 81 | self.reset_parameters() 82 | 83 | def reset_parameters(self): 84 | stdv = 1. / math.sqrt(self.d) 85 | self.A.data.uniform_(-stdv, stdv) 86 | 87 | def forward(self, X, Y): 88 | assert X.shape[2] == Y.shape[2] == self.d 89 | X = X.unsqueeze(1) 90 | Y = Y.unsqueeze(2) 91 | dxy = X - Y 92 | M = torch.matmul(self.A, self.A.transpose(0, 1)) 93 | M = torch.matmul(dxy.unsqueeze(-2), M) 94 | M = torch.matmul(M, dxy.unsqueeze(-1)) 95 | 96 | return self.relu(M.squeeze()) 97 | 98 | 99 | class AffinityFC(nn.Module): 100 | """ 101 | Affinity Layer to compute the affinity matrix from feature space. 102 | Affinity score is modeled by a fc neural network. 103 | Parameter: input dimension d, list of hidden layer dimension hds 104 | Input: feature X, Y 105 | Output: affinity matrix M 106 | """ 107 | def __init__(self, d, hds=None): 108 | super(AffinityFC, self).__init__() 109 | self.d = d 110 | if hds is None: 111 | self.hds = [1024,] 112 | else: 113 | self.hds = hds 114 | self.hds.append(1) 115 | fc_lst = [] 116 | last_hd = self.d * 2 117 | for hd in self.hds: 118 | fc_lst.append(nn.Linear(last_hd, hd)) 119 | fc_lst.append(nn.ReLU()) 120 | last_hd = hd 121 | 122 | self.fc = nn.Sequential(*fc_lst[:-1]) # last relu omitted 123 | 124 | def forward(self, X, Y): 125 | assert X.shape[2] == Y.shape[2] == self.d 126 | cat_feat = torch.cat((X.unsqueeze(-2).expand(X.shape[0], X.shape[1], Y.shape[1], X.shape[2]), 127 | Y.unsqueeze(-3).expand(Y.shape[0], X.shape[1], Y.shape[1], Y.shape[2])), dim=-1) 128 | result = self.fc(cat_feat).squeeze(-1) 129 | return result 130 | 131 | 132 | class AffinityBiFC(nn.Module): 133 | """ 134 | Affinity Layer to compute the affinity matrix from feature space. 135 | Affinity score is modeled by a bilinear layer followed by a fc neural network. 136 | Parameter: input dimension d, biliear dimension bd, list of hidden layer dimension hds 137 | Input: feature X, Y 138 | Output: affinity matrix M 139 | """ 140 | def __init__(self, d, bd=1024, hds=None): 141 | super(AffinityBiFC, self).__init__() 142 | self.d = d 143 | self.bd = bd 144 | if hds is None: 145 | self.hds = [] 146 | self.hds.append(1) 147 | 148 | self.A = Parameter(Tensor(self.d, self.d, self.bd)) 149 | self.reset_parameters() 150 | 151 | fc_lst = [] 152 | last_hd = self.bd 153 | for hd in self.hds: 154 | fc_lst.append(nn.Linear(last_hd, hd)) 155 | fc_lst.append(nn.ReLU()) 156 | last_hd = hd 157 | self.fc = nn.Sequential(*fc_lst[:-1]) # last relu omitted 158 | 159 | def reset_parameters(self): 160 | stdv = 1. / math.sqrt(self.d) 161 | self.A.data.uniform_(-stdv, stdv) 162 | 163 | def forward(self, X, Y): 164 | device = X.device 165 | assert X.shape[2] == Y.shape[2] == self.d 166 | bi_result = torch.empty(X.shape[0], X.shape [1], Y.shape[1], self.bd, device=device) 167 | for i in range(self.bd): 168 | tmp = torch.matmul(X, self.A[:, :, i]) 169 | tmp = torch.matmul(tmp, Y.transpose(1, 2)) 170 | bi_result[:, :, :, i] = tmp 171 | S = self.fc(bi_result) 172 | assert len(S.shape) == 3 173 | return S 174 | -------------------------------------------------------------------------------- /models/CIE/model.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | from src.lap_solvers.sinkhorn import Sinkhorn 5 | from src.feature_align import feature_align 6 | from src.gconv import Siamese_ChannelIndependentConv #, Siamese_GconvEdgeDPP, Siamese_GconvEdgeOri 7 | from models.PCA.affinity_layer import Affinity 8 | from src.lap_solvers.hungarian import hungarian 9 | 10 | from src.utils.config import cfg 11 | 12 | from src.backbone import * 13 | CNN = eval(cfg.BACKBONE) 14 | 15 | 16 | class Net(CNN): 17 | def __init__(self): 18 | super(Net, self).__init__() 19 | self.sinkhorn = Sinkhorn(max_iter=cfg.CIE.SK_ITER_NUM, epsilon=cfg.CIE.SK_EPSILON, tau=cfg.CIE.SK_TAU) 20 | self.l2norm = nn.LocalResponseNorm(cfg.CIE.FEATURE_CHANNEL * 2, alpha=cfg.CIE.FEATURE_CHANNEL * 2, beta=0.5, k=0) 21 | self.gnn_layer = cfg.CIE.GNN_LAYER # numbur of GNN layers 22 | for i in range(self.gnn_layer): 23 | if i == 0: 24 | gnn_layer = Siamese_ChannelIndependentConv(cfg.CIE.FEATURE_CHANNEL * 2, cfg.CIE.GNN_FEAT, 1) 25 | else: 26 | gnn_layer = Siamese_ChannelIndependentConv(cfg.CIE.GNN_FEAT, cfg.CIE.GNN_FEAT, cfg.CIE.GNN_FEAT) 27 | self.add_module('gnn_layer_{}'.format(i), gnn_layer) 28 | self.add_module('affinity_{}'.format(i), Affinity(cfg.CIE.GNN_FEAT)) 29 | if i == self.gnn_layer - 2: # only second last layer will have cross-graph module 30 | self.add_module('cross_graph_{}'.format(i), nn.Linear(cfg.CIE.GNN_FEAT * 2, cfg.CIE.GNN_FEAT)) 31 | self.add_module('cross_graph_edge_{}'.format(i), nn.Linear(cfg.CIE.GNN_FEAT * 2, cfg.CIE.GNN_FEAT)) 32 | self.rescale = cfg.PROBLEM.RESCALE 33 | 34 | def forward(self, data_dict, **kwargs): 35 | if 'images' in data_dict: 36 | # real image data 37 | src, tgt = data_dict['images'] 38 | P_src, P_tgt = data_dict['Ps'] 39 | ns_src, ns_tgt = data_dict['ns'] 40 | G_src, G_tgt = data_dict['Gs'] 41 | H_src, H_tgt = data_dict['Hs'] 42 | # extract feature 43 | src_node = self.node_layers(src) 44 | src_edge = self.edge_layers(src_node) 45 | tgt_node = self.node_layers(tgt) 46 | tgt_edge = self.edge_layers(tgt_node) 47 | 48 | # feature normalization 49 | src_node = self.l2norm(src_node) 50 | src_edge = self.l2norm(src_edge) 51 | tgt_node = self.l2norm(tgt_node) 52 | tgt_edge = self.l2norm(tgt_edge) 53 | 54 | # arrange features 55 | U_src = feature_align(src_node, P_src, ns_src, self.rescale) 56 | F_src = feature_align(src_edge, P_src, ns_src, self.rescale) 57 | U_tgt = feature_align(tgt_node, P_tgt, ns_tgt, self.rescale) 58 | F_tgt = feature_align(tgt_edge, P_tgt, ns_tgt, self.rescale) 59 | elif 'features' in data_dict: 60 | # synthetic data 61 | src, tgt = data_dict['features'] 62 | ns_src, ns_tgt = data_dict['ns'] 63 | G_src, G_tgt = data_dict['Gs'] 64 | H_src, H_tgt = data_dict['Hs'] 65 | 66 | U_src = src[:, :src.shape[1] // 2, :] 67 | F_src = src[:, src.shape[1] // 2:, :] 68 | U_tgt = tgt[:, :tgt.shape[1] // 2, :] 69 | F_tgt = tgt[:, tgt.shape[1] // 2:, :] 70 | else: 71 | raise ValueError('Unknown data type for this model.') 72 | 73 | P_src_dis = (P_src.unsqueeze(1) - P_src.unsqueeze(2)) 74 | P_src_dis = torch.norm(P_src_dis, p=2, dim=3).detach() 75 | P_tgt_dis = (P_tgt.unsqueeze(1) - P_tgt.unsqueeze(2)) 76 | P_tgt_dis = torch.norm(P_tgt_dis, p=2, dim=3).detach() 77 | 78 | Q_src = torch.exp(-P_src_dis / self.rescale[0]) 79 | Q_tgt = torch.exp(-P_tgt_dis / self.rescale[0]) 80 | 81 | emb_edge1 = Q_src.unsqueeze(-1) 82 | emb_edge2 = Q_tgt.unsqueeze(-1) 83 | 84 | # adjacency matrices 85 | A_src = torch.bmm(G_src, H_src.transpose(1, 2)) 86 | A_tgt = torch.bmm(G_tgt, H_tgt.transpose(1, 2)) 87 | 88 | # U_src, F_src are features at different scales 89 | emb1, emb2 = torch.cat((U_src, F_src), dim=1).transpose(1, 2), torch.cat((U_tgt, F_tgt), dim=1).transpose(1, 2) 90 | ss = [] 91 | 92 | for i in range(self.gnn_layer): 93 | gnn_layer = getattr(self, 'gnn_layer_{}'.format(i)) 94 | 95 | # during forward process, the network structure will not change 96 | emb1, emb2, emb_edge1, emb_edge2 = gnn_layer([A_src, emb1, emb_edge1], [A_tgt, emb2, emb_edge2]) 97 | 98 | affinity = getattr(self, 'affinity_{}'.format(i)) 99 | s = affinity(emb1, emb2) # xAx^T 100 | 101 | s = self.sinkhorn(s, ns_src, ns_tgt) 102 | ss.append(s) 103 | 104 | if i == self.gnn_layer - 2: 105 | cross_graph = getattr(self, 'cross_graph_{}'.format(i)) 106 | new_emb1 = cross_graph(torch.cat((emb1, torch.bmm(s, emb2)), dim=-1)) 107 | new_emb2 = cross_graph(torch.cat((emb2, torch.bmm(s.transpose(1, 2), emb1)), dim=-1)) 108 | emb1 = new_emb1 109 | emb2 = new_emb2 110 | 111 | # edge cross embedding 112 | ''' 113 | cross_graph_edge = getattr(self, 'cross_graph_edge_{}'.format(i)) 114 | emb_edge1 = emb_edge1.permute(0, 3, 1, 2) 115 | emb_edge2 = emb_edge2.permute(0, 3, 1, 2) 116 | s = s.unsqueeze(1) 117 | new_emb_edge1 = cross_graph_edge(torch.cat((emb_edge1, torch.matmul(torch.matmul(s, emb_edge2), s.transpose(2, 3))), dim=1).permute(0, 2, 3, 1)) 118 | new_emb_edge2 = cross_graph_edge(torch.cat((emb_edge2, torch.matmul(torch.matmul(s.transpose(2, 3), emb_edge1), s)), dim=1).permute(0, 2, 3, 1)) 119 | emb_edge1 = new_emb_edge1 120 | emb_edge2 = new_emb_edge2 121 | ''' 122 | 123 | data_dict.update({ 124 | 'ds_mat': ss[-1], 125 | 'perm_mat': hungarian(ss[-1], ns_src, ns_tgt) 126 | }) 127 | return data_dict -------------------------------------------------------------------------------- /src/dataset/qaplib.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from src.utils.config import cfg 3 | from pathlib import Path 4 | from src.dataset.base_dataset import BaseDataset 5 | import os 6 | import re 7 | import urllib 8 | 9 | 10 | cls_list = ['bur', 'chr', 'els', 'esc', 'had', 'kra', 'lipa', 'nug', 'rou', 'scr', 'sko', 'ste', 'tai', 'tho', 'wil'] 11 | 12 | class QAPLIB(BaseDataset): 13 | def __init__(self, sets, cls, fetch_online=False): 14 | super(QAPLIB, self).__init__() 15 | self.classes = ['qaplib'] 16 | self.sets = sets 17 | 18 | if cls is not None and cls != 'none': 19 | idx = cls_list.index(cls) 20 | self.cls_list = [cls_list[idx]] 21 | else: 22 | self.cls_list = cls_list 23 | 24 | self.data_list = [] 25 | # if to fetch data online, need first download, then ... 26 | self.qap_path = Path(cfg.QAPLIB.DIR) 27 | # import pdb; pdb.set_trace() 28 | 29 | for inst in self.cls_list: 30 | for dat_path in self.qap_path.glob(inst + '*.dat'): 31 | name = dat_path.name[:-4] 32 | prob_size = int(re.findall(r"\d+", name)[0]) 33 | if (self.sets == 'test' and prob_size > cfg.QAPLIB.MAX_TEST_SIZE) \ 34 | or (self.sets == 'train' and prob_size > cfg.QAPLIB.MAX_TRAIN_SIZE): 35 | continue 36 | # TODO: check if the solution file exits 37 | if not os.path.exists(self.qap_path / (name + '.sln')): 38 | continue 39 | self.data_list.append(name) 40 | 41 | # remove trivial instance esc16f 42 | if 'esc16f' in self.data_list: 43 | self.data_list.remove('esc16f') 44 | 45 | # define compare function 46 | def name_cmp(a, b): 47 | a = re.findall(r'[0-9]+|[a-z]+', a) 48 | b = re.findall(r'[0-9]+|[a-z]+', b) 49 | for _a, _b in zip(a, b): 50 | if _a.isdigit() and _b.isdigit(): 51 | _a = int(_a) 52 | _b = int(_b) 53 | cmp = (_a > _b) - (_a < _b) 54 | if cmp != 0: 55 | return cmp 56 | if len(a) > len(b): 57 | return -1 58 | elif len(a) < len(b): 59 | return 1 60 | else: 61 | return 0 62 | 63 | def cmp_to_key(mycmp): 64 | 'Convert a cmp= function into a key= function' 65 | class K: 66 | def __init__(self, obj, *args): 67 | self.obj = obj 68 | def __lt__(self, other): 69 | return mycmp(self.obj, other.obj) < 0 70 | def __gt__(self, other): 71 | return mycmp(self.obj, other.obj) > 0 72 | def __eq__(self, other): 73 | return mycmp(self.obj, other.obj) == 0 74 | def __le__(self, other): 75 | return mycmp(self.obj, other.obj) <= 0 76 | def __ge__(self, other): 77 | return mycmp(self.obj, other.obj) >= 0 78 | def __ne__(self, other): 79 | return mycmp(self.obj, other.obj) != 0 80 | return K 81 | 82 | # sort data list according to the names 83 | self.data_list.sort(key=cmp_to_key(name_cmp)) 84 | 85 | fetched_flag = self.qap_path / 'fetched_online' 86 | 87 | if fetch_online or not fetched_flag.exists(): 88 | self.__fetch_online() 89 | fetched_flag.touch() 90 | 91 | def get_pair(self, idx, shuffle=None): 92 | """ 93 | Get QAP data by index 94 | :param idx: dataset index 95 | :param shuffle: no use here 96 | :return: (pair of data, groundtruth permutation matrix) 97 | """ 98 | name = self.data_list[idx] 99 | 100 | dat_path = self.qap_path / (name + '.dat') 101 | sln_path = self.qap_path / (name + '.sln') 102 | dat_file = dat_path.open() 103 | sln_file = sln_path.open() 104 | 105 | def split_line(x): 106 | for _ in re.split(r'[,\s]', x.rstrip('\n')): 107 | if _ == "": 108 | continue 109 | else: 110 | yield int(_) 111 | dat_list = [] 112 | for line in dat_file: 113 | if line == '\n': continue 114 | dat_list.append([_ for _ in split_line(line)]) 115 | # dat_list = [[_ for _ in split_line(line)] for line in dat_file] 116 | sln_list = [[_ for _ in split_line(line)] for line in sln_file] 117 | # print(dat_list) 118 | # print('Current Name', name) 119 | # if name == 'lipa40a': 120 | # import pdb; pdb.set_trace(); 121 | prob_size = dat_list[0][0] 122 | 123 | # read data 124 | r = 0 125 | c = 0 126 | Fi = [[]] 127 | Fj = [[]] 128 | F = Fi 129 | for l in dat_list[1:]: 130 | F[r] += l 131 | c += len(l) 132 | assert c <= prob_size 133 | if c == prob_size: 134 | r += 1 135 | if r < prob_size: 136 | F.append([]) 137 | c = 0 138 | else: 139 | F = Fj 140 | r = 0 141 | c = 0 142 | Fi = np.array(Fi, dtype=np.float32) 143 | Fj = np.array(Fj, dtype=np.float32) 144 | assert Fi.shape == Fj.shape == (prob_size, prob_size) 145 | #K = np.kron(Fj, Fi) 146 | 147 | # read solution 148 | sol = sln_list[0][1] 149 | perm_list = [] 150 | for _ in sln_list[1:]: 151 | perm_list += _ 152 | assert len(perm_list) == prob_size 153 | perm_mat = np.zeros((prob_size, prob_size), dtype=np.float32) 154 | for r, c in enumerate(perm_list): 155 | perm_mat[r, c - 1] = 1 156 | 157 | return Fi, Fj, perm_mat, sol, name 158 | 159 | def __fetch_online(self): 160 | """ 161 | Fetch from online QAPLIB data 162 | """ 163 | for name in self.data_list: 164 | dat_content = urllib.request.urlopen(cfg.QAPLIB.ONLINE_REPO + 'data.d/{}.dat'.format(name)).read() 165 | sln_content = urllib.request.urlopen(cfg.QAPLIB.ONLINE_REPO + 'soln.d/{}.sln'.format(name)).read() 166 | 167 | dat_file = (self.qap_path / (name + '.dat')).open('wb') 168 | dat_file.write(dat_content) 169 | sln_file = (self.qap_path / (name + '.sln')).open('wb') 170 | sln_file.write(sln_content) 171 | -------------------------------------------------------------------------------- /src/dataset/imc_pt_sparsegm.py: -------------------------------------------------------------------------------- 1 | from pathlib import Path 2 | from PIL import Image 3 | import numpy as np 4 | from src.utils.config import cfg 5 | from src.dataset.base_dataset import BaseDataset 6 | import random 7 | 8 | 9 | class IMC_PT_SparseGM(BaseDataset): 10 | def __init__(self, sets, obj_resize): 11 | """ 12 | :param sets: 'train' or 'test' 13 | :param obj_resize: resized object size 14 | """ 15 | super(IMC_PT_SparseGM, self).__init__() 16 | assert sets in ('train', 'test'), 'No match found for dataset {}'.format(sets) 17 | self.sets = sets 18 | self.classes = cfg.IMC_PT_SparseGM.CLASSES[sets] 19 | self.total_kpt_num = cfg.IMC_PT_SparseGM.TOTAL_KPT_NUM 20 | 21 | self.root_path_npz = Path(cfg.IMC_PT_SparseGM.ROOT_DIR_NPZ) 22 | self.root_path_img = Path(cfg.IMC_PT_SparseGM.ROOT_DIR_IMG) 23 | self.obj_resize = obj_resize 24 | 25 | self.img_lists = [np.load(self.root_path_npz / cls / 'img_info.npz')['img_name'].tolist() 26 | for cls in self.classes] 27 | 28 | def get_pair(self, cls=None, shuffle=True, tgt_outlier=False, src_outlier=False): 29 | """ 30 | Randomly get a pair of objects from Photo Tourism dataset 31 | :param cls: None for random class, or specify for a certain set 32 | :param shuffle: random shuffle the keypoints 33 | :param src_outlier: allow outlier in the source graph (first graph) 34 | :param tgt_outlier: allow outlier in the target graph (second graph) 35 | :return: (pair of data, groundtruth permutation matrix) 36 | """ 37 | if cls is None: 38 | cls = random.randrange(0, len(self.classes)) 39 | elif type(cls) == str: 40 | cls = self.classes.index(cls) 41 | assert type(cls) == int and 0 <= cls < len(self.classes) 42 | 43 | anno_pair = [] 44 | for img_name in random.sample(self.img_lists[cls], 2): 45 | anno_dict = self.__get_anno_dict(img_name, cls) 46 | if shuffle: 47 | random.shuffle(anno_dict['keypoints']) 48 | anno_pair.append(anno_dict) 49 | 50 | perm_mat = np.zeros([len(_['keypoints']) for _ in anno_pair], dtype=np.float32) 51 | row_list = [] 52 | col_list = [] 53 | for i, keypoint in enumerate(anno_pair[0]['keypoints']): 54 | for j, _keypoint in enumerate(anno_pair[1]['keypoints']): 55 | if keypoint['name'] == _keypoint['name']: 56 | if keypoint['name'] != 'outlier': 57 | perm_mat[i, j] = 1 58 | row_list.append(i) 59 | col_list.append(j) 60 | break 61 | row_list.sort() 62 | col_list.sort() 63 | if not src_outlier: 64 | perm_mat = perm_mat[row_list, :] 65 | anno_pair[0]['keypoints'] = [anno_pair[0]['keypoints'][i] for i in row_list] 66 | if not tgt_outlier: 67 | perm_mat = perm_mat[:, col_list] 68 | anno_pair[1]['keypoints'] = [anno_pair[1]['keypoints'][j] for j in col_list] 69 | 70 | return anno_pair, perm_mat 71 | 72 | def get_multi(self, cls=None, num=2, shuffle=True, filter_outlier=False): 73 | """ 74 | Randomly get multiple objects from Willow Object Class dataset for multi-matching. 75 | :param cls: None for random class, or specify for a certain set 76 | :param num: number of objects to be fetched 77 | :param shuffle: random shuffle the keypoints 78 | :param filter_outlier: filter out outlier keypoints among images 79 | :return: (list of data, list of permutation matrices) 80 | """ 81 | assert not filter_outlier, 'Multi-matching on IMC_PT_SparseGM dataset with filtered outliers is not supported' 82 | 83 | if cls is None: 84 | cls = random.randrange(0, len(self.classes)) 85 | elif type(cls) == str: 86 | cls = self.classes.index(cls) 87 | assert type(cls) == int and 0 <= cls < len(self.classes) 88 | 89 | anno_list = [] 90 | for img_name in random.sample(self.img_lists[cls], num): 91 | anno_dict = self.__get_anno_dict(img_name, cls) 92 | if shuffle: 93 | random.shuffle(anno_dict['keypoints']) 94 | anno_list.append(anno_dict) 95 | 96 | perm_mat = [np.zeros([self.total_kpt_num, len(x['keypoints'])], dtype=np.float32) for x in anno_list] 97 | for k, anno_dict in enumerate(anno_list): 98 | kpt_name_list = [x['name'] for x in anno_dict['keypoints']] 99 | for j, name in enumerate(kpt_name_list): 100 | perm_mat[k][name, j] = 1 101 | for k in range(num): 102 | perm_mat[k] = perm_mat[k].transpose() 103 | 104 | return anno_list, perm_mat 105 | 106 | def __get_anno_dict(self, img_name, cls): 107 | """ 108 | Get an annotation dict from .npz annotation 109 | """ 110 | img_file = self.root_path_img / self.classes[cls] / 'dense' / 'images' / img_name 111 | npz_file = self.root_path_npz / self.classes[cls] / (img_name + '.npz') 112 | 113 | assert img_file.exists(), '{} does not exist.'.format(img_file) 114 | assert npz_file.exists(), '{} does not exist.'.format(npz_file) 115 | 116 | with Image.open(str(img_file)) as img: 117 | ori_sizes = img.size 118 | obj = img.resize(self.obj_resize, resample=Image.BICUBIC) 119 | xmin = 0 120 | ymin = 0 121 | w = ori_sizes[0] 122 | h = ori_sizes[1] 123 | 124 | with np.load(str(npz_file)) as npz_anno: 125 | kpts = npz_anno['points'] 126 | if len(kpts.shape) != 2: 127 | ValueError('{} contains no keypoints.'.format(img_file)) 128 | 129 | keypoint_list = [] 130 | for i in range(kpts.shape[1]): 131 | kpt_index = int(kpts[0, i]) 132 | assert kpt_index < self.total_kpt_num 133 | attr = { 134 | 'name': kpt_index, 135 | 'x': kpts[1, i] * self.obj_resize[0] / w, 136 | 'y': kpts[2, i] * self.obj_resize[1] / h 137 | } 138 | keypoint_list.append(attr) 139 | 140 | anno_dict = dict() 141 | anno_dict['image'] = obj 142 | anno_dict['keypoints'] = keypoint_list 143 | anno_dict['bounds'] = xmin, ymin, w, h 144 | anno_dict['ori_sizes'] = ori_sizes 145 | anno_dict['cls'] = self.classes[cls] 146 | anno_dict['univ_size'] = self.total_kpt_num 147 | 148 | return anno_dict 149 | 150 | def len(self, cls): 151 | if type(cls) == int: 152 | cls = self.classes[cls] 153 | assert cls in self.classes 154 | return len(self.img_lists[self.classes.index(cls)]) 155 | -------------------------------------------------------------------------------- /models/BBGM/model.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import itertools 3 | 4 | from models.BBGM.affinity_layer import InnerProductWithWeightsAffinity 5 | from models.BBGM.sconv_archs import SiameseSConvOnNodes, SiameseNodeFeaturesToEdgeFeatures 6 | from lpmp_py import GraphMatchingModule 7 | from lpmp_py import MultiGraphMatchingModule 8 | from src.feature_align import feature_align 9 | 10 | from src.utils.config import cfg 11 | 12 | from src.backbone import * 13 | CNN = eval(cfg.BACKBONE) 14 | 15 | 16 | def lexico_iter(lex): 17 | return itertools.combinations(lex, 2) 18 | 19 | def normalize_over_channels(x): 20 | channel_norms = torch.norm(x, dim=1, keepdim=True) 21 | return x / channel_norms 22 | 23 | 24 | def concat_features(embeddings, num_vertices): 25 | res = torch.cat([embedding[:, :num_v] for embedding, num_v in zip(embeddings, num_vertices)], dim=-1) 26 | return res.transpose(0, 1) 27 | 28 | 29 | class Net(CNN): 30 | def __init__(self): 31 | super(Net, self).__init__() 32 | self.message_pass_node_features = SiameseSConvOnNodes(input_node_dim=cfg.BBGM.FEATURE_CHANNEL) 33 | self.build_edge_features_from_node_features = SiameseNodeFeaturesToEdgeFeatures( 34 | total_num_nodes=self.message_pass_node_features.num_node_features 35 | ) 36 | self.global_state_dim = cfg.BBGM.FEATURE_CHANNEL 37 | self.vertex_affinity = InnerProductWithWeightsAffinity( 38 | self.global_state_dim, self.message_pass_node_features.num_node_features) 39 | self.edge_affinity = InnerProductWithWeightsAffinity( 40 | self.global_state_dim, 41 | self.build_edge_features_from_node_features.num_edge_features) 42 | self.rescale = cfg.PROBLEM.RESCALE 43 | 44 | def forward( 45 | self, 46 | data_dict, 47 | ): 48 | images = data_dict['images'] 49 | points = data_dict['Ps'] 50 | n_points = data_dict['ns'] 51 | graphs = data_dict['pyg_graphs'] 52 | num_graphs = len(images) 53 | 54 | if cfg.PROBLEM.TYPE == '2GM' and 'gt_perm_mat' in data_dict: 55 | gt_perm_mats = [data_dict['gt_perm_mat']] 56 | elif cfg.PROBLEM.TYPE == 'MGM' and 'gt_perm_mat' in data_dict: 57 | perm_mat_list = data_dict['gt_perm_mat'] 58 | gt_perm_mats = [torch.bmm(pm_src, pm_tgt.transpose(1, 2)) for pm_src, pm_tgt in lexico_iter(perm_mat_list)] 59 | else: 60 | raise ValueError('Ground truth information is required during training.') 61 | 62 | global_list = [] 63 | orig_graph_list = [] 64 | for image, p, n_p, graph in zip(images, points, n_points, graphs): 65 | # extract feature 66 | nodes = self.node_layers(image) 67 | edges = self.edge_layers(nodes) 68 | 69 | global_list.append(self.final_layers(edges).reshape((nodes.shape[0], -1))) 70 | nodes = normalize_over_channels(nodes) 71 | edges = normalize_over_channels(edges) 72 | 73 | # arrange features 74 | U = concat_features(feature_align(nodes, p, n_p, self.rescale), n_p) 75 | F = concat_features(feature_align(edges, p, n_p, self.rescale), n_p) 76 | node_features = torch.cat((U, F), dim=1) 77 | graph.x = node_features 78 | 79 | graph = self.message_pass_node_features(graph) 80 | orig_graph = self.build_edge_features_from_node_features(graph) 81 | orig_graph_list.append(orig_graph) 82 | 83 | global_weights_list = [ 84 | torch.cat([global_src, global_tgt], axis=-1) for global_src, global_tgt in lexico_iter(global_list) 85 | ] 86 | 87 | global_weights_list = [normalize_over_channels(g) for g in global_weights_list] 88 | 89 | unary_costs_list = [ 90 | self.vertex_affinity([item.x for item in g_1], [item.x for item in g_2], global_weights) 91 | for (g_1, g_2), global_weights in zip(lexico_iter(orig_graph_list), global_weights_list) 92 | ] 93 | 94 | # Similarities to costs 95 | unary_costs_list = [[-x for x in unary_costs] for unary_costs in unary_costs_list] 96 | 97 | if self.training: 98 | unary_costs_list = [ 99 | [ 100 | x + 1.0*gt[:dim_src, :dim_tgt] # Add margin with alpha = 1.0 101 | for x, gt, dim_src, dim_tgt in zip(unary_costs, gt_perm_mat, ns_src, ns_tgt) 102 | ] 103 | for unary_costs, gt_perm_mat, (ns_src, ns_tgt) in zip(unary_costs_list, gt_perm_mats, lexico_iter(n_points)) 104 | ] 105 | 106 | quadratic_costs_list = [ 107 | self.edge_affinity([item.edge_attr for item in g_1], [item.edge_attr for item in g_2], global_weights) 108 | for (g_1, g_2), global_weights in zip(lexico_iter(orig_graph_list), global_weights_list) 109 | ] 110 | 111 | # Similarities to costs 112 | quadratic_costs_list = [[-0.5 * x for x in quadratic_costs] for quadratic_costs in quadratic_costs_list] 113 | 114 | if cfg.BBGM.SOLVER_NAME == "LPMP": 115 | all_edges = [[item.edge_index for item in graph] for graph in orig_graph_list] 116 | gm_solvers = [ 117 | GraphMatchingModule( 118 | all_left_edges, 119 | all_right_edges, 120 | ns_src, 121 | ns_tgt, 122 | cfg.BBGM.LAMBDA_VAL, 123 | cfg.BBGM.SOLVER_PARAMS, 124 | ) 125 | for (all_left_edges, all_right_edges), (ns_src, ns_tgt) in zip( 126 | lexico_iter(all_edges), lexico_iter(n_points) 127 | ) 128 | ] 129 | matchings = [ 130 | gm_solver(unary_costs, quadratic_costs) 131 | for gm_solver, unary_costs, quadratic_costs in zip(gm_solvers, unary_costs_list, quadratic_costs_list) 132 | ] 133 | elif cfg.BBGM.SOLVER_NAME == "LPMP_MGM": 134 | all_edges = [[item.edge_index for item in graph] for graph in orig_graph_list] 135 | gm_solver = MultiGraphMatchingModule( 136 | all_edges, n_points, cfg.BBGM.LAMBDA_VAL, cfg.BBGM.SOLVER_PARAMS) 137 | matchings = gm_solver(unary_costs_list, quadratic_costs_list) 138 | else: 139 | raise ValueError("Unknown solver {}".format(cfg.BBGM.SOLVER_NAME)) 140 | 141 | 142 | if cfg.PROBLEM.TYPE == '2GM': 143 | data_dict.update({ 144 | 'ds_mat': None, 145 | 'perm_mat': matchings[0] 146 | }) 147 | elif cfg.PROBLEM.TYPE == 'MGM': 148 | indices = list(lexico_iter(range(num_graphs))) 149 | data_dict.update({ 150 | 'perm_mat_list': matchings, 151 | 'graph_indices': indices, 152 | 'gt_perm_mat_list': gt_perm_mats 153 | }) 154 | 155 | return data_dict 156 | -------------------------------------------------------------------------------- /src/factorize_graph_matching.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import Tensor 3 | from torch.autograd import Function 4 | from src.utils.sparse import bilinear_diag_torch 5 | from src.sparse_torch import CSRMatrix3d, CSCMatrix3d 6 | import scipy.sparse as ssp 7 | import numpy as np 8 | 9 | 10 | def construct_aff_mat(Ke: Tensor, Kp: Tensor, KroG: CSRMatrix3d, KroH: CSCMatrix3d, 11 | KroGt: CSRMatrix3d=None, KroHt: CSCMatrix3d=None) -> Tensor: 12 | r""" 13 | Construct the complete affinity matrix with edge-wise affinity matrix :math:`\mathbf{K}_e`, node-wise matrix 14 | :math:`\mathbf{K}_p` and graph connectivity matrices :math:`\mathbf{G}_1, \mathbf{H}_1, \mathbf{G}_2, \mathbf{H}_2` 15 | 16 | .. math :: 17 | \mathbf{K}=\mathrm{diag}(\mathrm{vec}(\mathbf{K}_p)) + 18 | (\mathbf{G}_2 \otimes_{\mathcal{K}} \mathbf{G}_1) \mathrm{diag}(\mathrm{vec}(\mathbf{K}_e)) 19 | (\mathbf{H}_2 \otimes_{\mathcal{K}} \mathbf{H}_1)^\top 20 | 21 | where :math:`\mathrm{diag}(\cdot)` means building a diagonal matrix based on the given vector, 22 | and :math:`\mathrm{vec}(\cdot)` means column-wise vectorization. 23 | :math:`\otimes_{\mathcal{K}}` denotes Kronecker product. 24 | 25 | This function supports batched operations. This formulation is developed by `"F. Zhou and F. Torre. Factorized 26 | Graph Matching. TPAMI 2015." `_ 27 | 28 | :param Ke: :math:`(b\times n_{e_1}\times n_{e_2})` edge-wise affinity matrix. 29 | :math:`n_{e_1}`: number of edges in graph 1, :math:`n_{e_2}`: number of edges in graph 2 30 | :param Kp: :math:`(b\times n_1\times n_2)` node-wise affinity matrix. 31 | :math:`n_1`: number of nodes in graph 1, :math:`n_2`: number of nodes in graph 2 32 | :param KroG: :math:`(b\times n_1n_2 \times n_{e_1}n_{e_2})` kronecker product of 33 | :math:`\mathbf{G}_2 (b\times n_2 \times n_{e_2})`, :math:`\mathbf{G}_1 (b\times n_1 \times n_{e_1})` 34 | :param KroH: :math:`(b\times n_1n_2 \times n_{e_1}n_{e_2})` kronecker product of 35 | :math:`\mathbf{H}_2 (b\times n_2 \times n_{e_2})`, :math:`\mathbf{H}_1 (b\times n_1 \times n_{e_1})` 36 | :param KroGt: transpose of KroG (should be CSR, optional) 37 | :param KroHt: transpose of KroH (should be CSC, optional) 38 | :return: affinity matrix :math:`\mathbf K` 39 | 40 | .. note :: 41 | This function is optimized with sparse CSR and CSC matrices with GPU support for both forward and backward 42 | computation with PyTorch. To use this function, you need to install ``ninja-build``, ``gcc-7``, ``nvcc`` (which 43 | comes along with CUDA development tools) to successfully compile our customized CUDA code for CSR and CSC 44 | matrices. The compiler is automatically called upon requirement. 45 | 46 | For a graph matching problem with 5 nodes and 4 nodes, 47 | the connection of :math:`\mathbf K` and :math:`\mathbf{K}_p, \mathbf{K}_e` is illustrated as 48 | 49 | .. image :: ../../images/factorized_graph_matching.png 50 | 51 | where :math:`\mathbf K (20 \times 20)` is the complete affinity matrix, :math:`\mathbf{K}_p (5 \times 4)` is the 52 | node-wise affinity matrix, :math:`\mathbf{K}_e(16 \times 10)` is the edge-wise affinity matrix. 53 | """ 54 | return RebuildFGM.apply(Ke, Kp, KroG, KroH, KroGt, KroHt) 55 | 56 | 57 | def kronecker_torch(t1: Tensor, t2: Tensor) -> Tensor: 58 | r""" 59 | Compute the kronecker product of :math:`\mathbf{T}_1` and :math:`\mathbf{T}_2`. 60 | This function is implemented in torch API and is not efficient for sparse {0, 1} matrix. 61 | 62 | :param t1: input tensor 1 63 | :param t2: input tensor 2 64 | :return: kronecker product of :math:`\mathbf{T}_1` and :math:`\mathbf{T}_2` 65 | """ 66 | batch_num = t1.shape[0] 67 | t1dim1, t1dim2 = t1.shape[1], t1.shape[2] 68 | t2dim1, t2dim2 = t2.shape[1], t2.shape[2] 69 | if t1.is_sparse and t2.is_sparse: 70 | tt_idx = torch.stack(t1._indices()[0, :] * t2dim1, t1._indices()[1, :] * t2dim2) 71 | tt_idx = torch.repeat_interleave(tt_idx, t2._nnz(), dim=1) + t2._indices().repeat(1, t1._nnz()) 72 | tt_val = torch.repeat_interleave(t1._values(), t2._nnz(), dim=1) * t2._values().repeat(1, t1._nnz()) 73 | tt = torch.sparse.FloatTensor(tt_idx, tt_val, torch.Size(t1dim1 * t2dim1, t1dim2 * t2dim2)) 74 | else: 75 | t1 = t1.reshape(batch_num, -1, 1) 76 | t2 = t2.reshape(batch_num, 1, -1) 77 | tt = torch.bmm(t1, t2) 78 | tt = tt.reshape(batch_num, t1dim1, t1dim2, t2dim1, t2dim2) 79 | tt = tt.permute([0, 1, 3, 2, 4]) 80 | tt = tt.reshape(batch_num, t1dim1 * t2dim1, t1dim2 * t2dim2) 81 | return tt 82 | 83 | 84 | def kronecker_sparse(arr1: np.ndarray, arr2: np.ndarray) -> np.ndarray: 85 | r""" 86 | Compute the kronecker product of :math:`\mathbf{T}_1` and :math:`\mathbf{T}_2`. 87 | This function is implemented in scipy.sparse API and runs on cpu. 88 | 89 | :param arr1: input array 1 90 | :param arr2: input array 2 91 | :return: kronecker product of :math:`\mathbf{T}_1` and :math:`\mathbf{T}_2` 92 | """ 93 | s1 = ssp.coo_matrix(arr1) 94 | s2 = ssp.coo_matrix(arr2) 95 | ss = ssp.kron(s1, s2) 96 | return ss 97 | 98 | 99 | class RebuildFGM(Function): 100 | r""" 101 | Rebuild sparse affinity matrix in the formula of the paper `"Factorized Graph Matching, in 102 | TPAMI 2015" `_ 103 | 104 | See :func:`~src.factorize_graph_matching.construct_aff_mat` for detailed reference. 105 | """ 106 | @staticmethod 107 | def forward(ctx, Ke: Tensor, Kp: Tensor, Kro1: CSRMatrix3d, Kro2: CSCMatrix3d, 108 | Kro1T: CSRMatrix3d=None, Kro2T: CSCMatrix3d=None) -> Tensor: 109 | """ 110 | Forward function to compute the affinity matrix :math:`\mathbf K`. 111 | """ 112 | ctx.save_for_backward(Ke, Kp) 113 | if Kro1T is not None and Kro2T is not None: 114 | ctx.K = Kro1T, Kro2T 115 | else: 116 | ctx.K = Kro1.transpose(keep_type=True), Kro2.transpose(keep_type=True) 117 | batch_num = Ke.shape[0] 118 | 119 | Kro1Ke = Kro1.dotdiag(Ke.transpose(1, 2).contiguous().view(batch_num, -1)) 120 | Kro1KeKro2 = Kro1Ke.dot(Kro2, dense=True) 121 | 122 | K = torch.empty_like(Kro1KeKro2) 123 | for b in range(batch_num): 124 | K[b] = Kro1KeKro2[b] + torch.diag(Kp[b].transpose(0, 1).contiguous().view(-1)) 125 | 126 | return K 127 | 128 | @staticmethod 129 | def backward(ctx, dK): 130 | r""" 131 | Backward function from the affinity matrix :math:`\mathbf K` to node-wise affinity matrix :math:`\mathbf K_e` 132 | and edge-wize affinity matrix :math:`\mathbf K_e`. 133 | """ 134 | device = dK.device 135 | Ke, Kp = ctx.saved_tensors 136 | Kro1t, Kro2t = ctx.K 137 | dKe = dKp = None 138 | if ctx.needs_input_grad[0]: 139 | dKe = bilinear_diag_torch(Kro1t, dK.contiguous(), Kro2t) 140 | dKe = dKe.view(Ke.shape[0], Ke.shape[2], Ke.shape[1]).transpose(1, 2) 141 | if ctx.needs_input_grad[1]: 142 | dKp = torch.diagonal(dK, dim1=-2, dim2=-1) 143 | dKp = dKp.view(Kp.shape[0], Kp.shape[2], Kp.shape[1]).transpose(1, 2) 144 | 145 | return dKe, dKp, None, None, None, None 146 | -------------------------------------------------------------------------------- /src/gconv.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | from torch import Tensor 5 | from typing import Tuple, Optional, List, Union 6 | 7 | 8 | class Gconv(nn.Module): 9 | r""" 10 | Graph Convolutional Layer which is inspired and developed based on Graph Convolutional Network (GCN). 11 | Inspired by `Kipf and Welling. Semi-Supervised Classification with Graph Convolutional Networks. ICLR 2017. 12 | `_ 13 | 14 | :param in_features: the dimension of input node features 15 | :param out_features: the dimension of output node features 16 | """ 17 | def __init__(self, in_features: int, out_features: int): 18 | super(Gconv, self).__init__() 19 | self.num_inputs = in_features 20 | self.num_outputs = out_features 21 | self.a_fc = nn.Linear(self.num_inputs, self.num_outputs) 22 | self.u_fc = nn.Linear(self.num_inputs, self.num_outputs) 23 | 24 | def forward(self, A: Tensor, x: Tensor, norm: bool=True) -> Tensor: 25 | r""" 26 | Forward computation of graph convolution network. 27 | 28 | :param A: :math:`(b\times n\times n)` {0,1} adjacency matrix. :math:`b`: batch size, :math:`n`: number of nodes 29 | :param x: :math:`(b\times n\times d)` input node embedding. :math:`d`: feature dimension 30 | :param norm: normalize connectivity matrix or not 31 | :return: :math:`(b\times n\times d^\prime)` new node embedding 32 | """ 33 | if norm is True: 34 | A = F.normalize(A, p=1, dim=-2) 35 | ax = self.a_fc(x) 36 | ux = self.u_fc(x) 37 | x = torch.bmm(A, F.relu(ax)) + F.relu(ux) # has size (bs, N, num_outputs) 38 | return x 39 | 40 | 41 | class ChannelIndependentConv(nn.Module): 42 | r""" 43 | Channel Independent Embedding Convolution. 44 | Proposed by `"Yu et al. Learning deep graph matching with channel-independent embedding and Hungarian attention. 45 | ICLR 2020." `_ 46 | 47 | :param in_features: the dimension of input node features 48 | :param out_features: the dimension of output node features 49 | :param in_edges: the dimension of input edge features 50 | :param out_edges: (optional) the dimension of output edge features. It needs to be the same as ``out_features`` 51 | """ 52 | def __init__(self, in_features: int, out_features: int, in_edges: int, out_edges: int=None): 53 | super(ChannelIndependentConv, self).__init__() 54 | if out_edges is None: 55 | out_edges = out_features 56 | self.in_features = in_features 57 | self.out_features = out_features 58 | self.out_edges = out_edges 59 | # self.node_fc = nn.Linear(in_features, out_features // self.out_edges) 60 | self.node_fc = nn.Linear(in_features, out_features) 61 | self.node_sfc = nn.Linear(in_features, out_features) 62 | self.edge_fc = nn.Linear(in_edges, self.out_edges) 63 | 64 | def forward(self, A: Tensor, emb_node: Tensor, emb_edge: Tensor, mode: int=1) -> Tuple[Tensor, Tensor]: 65 | r""" 66 | :param A: :math:`(b\times n\times n)` {0,1} adjacency matrix. :math:`b`: batch size, :math:`n`: number of nodes 67 | :param emb_node: :math:`(b\times n\times d_n)` input node embedding. :math:`d_n`: node feature dimension 68 | :param emb_edge: :math:`(b\times n\times n\times d_e)` input edge embedding. :math:`d_e`: edge feature dimension 69 | :param mode: 1 or 2, refer to the paper for details 70 | :return: :math:`(b\times n\times d^\prime)` new node embedding, 71 | :math:`(b\times n\times n\times d^\prime)` new edge embedding 72 | """ 73 | if mode == 1: 74 | node_x = self.node_fc(emb_node) 75 | node_sx = self.node_sfc(emb_node) 76 | edge_x = self.edge_fc(emb_edge) 77 | 78 | A = A.unsqueeze(-1) 79 | A = torch.mul(A.expand_as(edge_x), edge_x) 80 | 81 | node_x = torch.matmul(A.transpose(2, 3).transpose(1, 2), 82 | node_x.unsqueeze(2).transpose(2, 3).transpose(1, 2)) 83 | node_x = node_x.squeeze(-1).transpose(1, 2) 84 | node_x = F.relu(node_x) + F.relu(node_sx) 85 | edge_x = F.relu(edge_x) 86 | 87 | return node_x, edge_x 88 | 89 | elif mode == 2: 90 | node_x = self.node_fc(emb_node) 91 | node_sx = self.node_sfc(emb_node) 92 | edge_x = self.edge_fc(emb_edge) 93 | 94 | d_x = node_x.unsqueeze(1) - node_x.unsqueeze(2) 95 | d_x = torch.sum(d_x ** 2, dim=3, keepdim=False) 96 | d_x = torch.exp(-d_x) 97 | 98 | A = A.unsqueeze(-1) 99 | A = torch.mul(A.expand_as(edge_x), edge_x) 100 | 101 | node_x = torch.matmul(A.transpose(2, 3).transpose(1, 2), 102 | node_x.unsqueeze(2).transpose(2, 3).transpose(1, 2)) 103 | node_x = node_x.squeeze(-1).transpose(1, 2) 104 | node_x = F.relu(node_x) + F.relu(node_sx) 105 | edge_x = F.relu(edge_x) 106 | return node_x, edge_x 107 | 108 | else: 109 | raise ValueError('Unknown mode {}. Possible options: 1 or 2'.format(mode)) 110 | 111 | 112 | class Siamese_Gconv(nn.Module): 113 | r""" 114 | Siamese Gconv neural network for processing arbitrary number of graphs. 115 | 116 | :param in_features: the dimension of input node features 117 | :param num_features: the dimension of output node features 118 | """ 119 | def __init__(self, in_features, num_features): 120 | super(Siamese_Gconv, self).__init__() 121 | self.gconv = Gconv(in_features, num_features) 122 | 123 | def forward(self, g1: Tuple[Tensor, Tensor, Tensor, int], *args) -> Union[Tensor, List[Tensor]]: 124 | r""" 125 | Forward computation of Siamese Gconv. 126 | 127 | :param g1: The first graph, which is a tuple of (:math:`(b\times n\times n)` {0,1} adjacency matrix, 128 | :math:`(b\times n\times d)` input node embedding, normalize connectivity matrix or not) 129 | :param args: Other graphs 130 | :return: A list of tensors composed of new node embeddings :math:`(b\times n\times d^\prime)` 131 | """ 132 | # embx are tensors of size (bs, N, num_features) 133 | emb1 = self.gconv(*g1) 134 | if len(args) == 0: 135 | return emb1 136 | else: 137 | returns = [emb1] 138 | for g in args: 139 | returns.append(self.gconv(*g)) 140 | return returns 141 | 142 | class Siamese_ChannelIndependentConv(nn.Module): 143 | r""" 144 | Siamese Channel Independent Conv neural network for processing arbitrary number of graphs. 145 | 146 | :param in_features: the dimension of input node features 147 | :param num_features: the dimension of output node features 148 | :param in_edges: the dimension of input edge features 149 | :param out_edges: (optional) the dimension of output edge features. It needs to be the same as ``num_features`` 150 | """ 151 | def __init__(self, in_features, num_features, in_edges, out_edges=None): 152 | super(Siamese_ChannelIndependentConv, self).__init__() 153 | self.in_feature = in_features 154 | self.gconv = ChannelIndependentConv(in_features, num_features, in_edges, out_edges) 155 | 156 | def forward(self, g1: Tuple[Tensor, Tensor, Optional[bool]], *args) -> List[Tensor]: 157 | r""" 158 | Forward computation of Siamese Channel Independent Conv. 159 | 160 | :param g1: The first graph, which is a tuple of (:math:`(b\times n\times n)` {0,1} adjacency matrix, 161 | :math:`(b\times n\times d_n)` input node embedding, :math:`(b\times n\times n\times d_e)` input edge embedding, 162 | mode (``1`` or ``2``)) 163 | :param args: Other graphs 164 | :return: A list of tensors composed of new node embeddings :math:`(b\times n\times d^\prime)`, appended with new 165 | edge embeddings :math:`(b\times n\times n\times d^\prime)` 166 | """ 167 | emb1, emb_edge1 = self.gconv(*g1) 168 | embs = [emb1] 169 | emb_edges = [emb_edge1] 170 | for g in args: 171 | emb2, emb_edge2 = self.gconv(*g) 172 | embs.append(emb2), emb_edges.append(emb_edge2) 173 | return embs + emb_edges 174 | -------------------------------------------------------------------------------- /src/utils/parse_args.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | from src.utils.config import cfg, cfg_from_file, cfg_from_list, get_output_dir, get_output_dir_new 3 | from pathlib import Path 4 | 5 | def parse_args(description): 6 | parser = argparse.ArgumentParser(description=description) 7 | ### Universal param 8 | parser.add_argument('--cfg', '--config', dest='cfg_file', action='append', 9 | help='an optional config file', default=None, type=str) 10 | parser.add_argument('--batch', dest='batch_size', 11 | help='batch size', default=None, type=int) 12 | ### Attack param 13 | parser.add_argument('--alpha', dest='alpha', 14 | help='alpha', default=None, type=float) 15 | parser.add_argument('--num_iter', dest='num_iter', 16 | help='iteration number', default=None, type=int) 17 | parser.add_argument('--warm_num_iter', default=None, type=int) 18 | parser.add_argument('--eps_feature', default=None, type=float) 19 | parser.add_argument('--eps_locality', default=None, type=float) 20 | parser.add_argument('--momentum_mu', default=0., type=float) 21 | 22 | parser.add_argument('--attack_type', default=None, type=str) 23 | parser.add_argument('--obj_type', default=None, type=str) 24 | parser.add_argument('--att_loss_func', default=None, type=str) 25 | parser.add_argument('--black',action='store_true', default=False) 26 | parser.add_argument('--inv', action='store_true', default=False) 27 | 28 | ### Model-concerning params 29 | parser.add_argument('--pretrained_path', default=None, type=str) 30 | 31 | ### Eval param 32 | parser.add_argument('--eval_num_iter', 33 | help='eval iteration number', default=None, type=int) 34 | parser.add_argument('--eval_alpha', 35 | help='eval step size per iteration', default=None, type=float) 36 | parser.add_argument('--eval_mode', default=None, type=str, choices=['clean', 'all', 'single', 'strong', 'supple', 'supple_s'], help='evaluation mode.') 37 | 38 | ### Train param 39 | parser.add_argument('--start_epoch', dest='start_epoch', 40 | help='start epoch number for resume training', default=None, type=int) 41 | parser.add_argument('--eval_epoch', default=None, type=int) 42 | parser.add_argument('--eval_per_num_epoch', default=None, type=int) 43 | 44 | parser.add_argument('--attack2_type', default=None, type=str) 45 | parser.add_argument('--obj2_type', default=None, type=str) 46 | parser.add_argument('--mode', default=None, type=str, choices=['at', '2step', 'eval'], help='training mode. pixel+loc: first perturb locations of points to generate new graph structure, then perturb features.') 47 | parser.add_argument('--burn_in_period', default=None, type=int) 48 | 49 | parser.add_argument('--sync_minmax', default=None, type=int, help='whether to use the same loss function or not.') 50 | parser.add_argument('--loss_func', default=None, type=str) 51 | 52 | parser.add_argument('--reg_level', default=None, type=int, help='regularization level') 53 | parser.add_argument('--reg_ratio', default=None, type=float) 54 | 55 | parser.add_argument('--affinity_mask', default=None, type=int) 56 | 57 | ### Utils params 58 | parser.add_argument("--exp_name", default=None, type=str) 59 | parser.add_argument("--prefix", default='', type=str, help='prefix to the experiment subdir name') 60 | parser.add_argument("--datetime", default=None, type=str) 61 | args = parser.parse_args() 62 | 63 | # load cfg from file 64 | if args.cfg_file is not None: 65 | for f in args.cfg_file: 66 | cfg_from_file(f) 67 | 68 | # load cfg from arguments 69 | if args.loss_func is not None: 70 | cfg_from_list(['TRAIN.LOSS_FUNC', args.loss_func]) 71 | if args.sync_minmax is not None: 72 | cfg_from_list(['TRAIN.SYNC_MINMAX', args.sync_minmax]) 73 | if args.affinity_mask is not None: 74 | cfg_from_list(['EVAL.AFF_MASK', args.affinity_mask]) 75 | if args.eval_mode is not None: 76 | cfg_from_list(['EVAL.MODE', args.eval_mode]) 77 | if args.batch_size is not None: 78 | cfg_from_list(['BATCH_SIZE', args.batch_size]) 79 | if args.start_epoch is not None: 80 | cfg_from_list(['TRAIN.START_EPOCH', args.start_epoch]) 81 | if args.eval_epoch is not None: 82 | cfg_from_list(['EVAL.EPOCH', args.eval_epoch]) 83 | if args.eval_per_num_epoch is not None: 84 | cfg_from_list(['EVAL.NUM_EPOCH', args.eval_per_num_epoch]) 85 | if args.alpha is not None: 86 | cfg_from_list(['ATTACK.ALPHA', args.alpha]) 87 | if args.eval_alpha is not None: 88 | cfg_from_list(['ATTACK.EVAL_ALPHA', args.eval_alpha]) 89 | if args.att_loss_func is not None: 90 | cfg_from_list(['ATTACK.LOSS_FUNC', args.att_loss_func]) 91 | if args.num_iter is not None: 92 | cfg_from_list(['ATTACK.STEP', args.num_iter]) 93 | if args.eval_num_iter is not None: 94 | cfg_from_list(['ATTACK.EVAL_STEP', args.eval_num_iter]) 95 | if args.eps_feature is not None: 96 | cfg_from_list(['ATTACK.EPSILON_FEATURE', args.eps_feature]) 97 | if args.eps_locality is not None: 98 | cfg_from_list(['ATTACK.EPSILON_LOCALITY', args.eps_locality]) 99 | if args.eps_feature is not None: 100 | cfg_from_list(['ATTACK.EVAL_EPSILON_FEATURE', args.eps_feature]) 101 | if args.eps_locality is not None: 102 | cfg_from_list(['ATTACK.EVAL_EPSILON_LOCALITY', args.eps_locality]) 103 | if args.momentum_mu is not None: 104 | cfg_from_list(['ATTACK.MU', args.momentum_mu]) 105 | if args.attack_type is not None: 106 | cfg_from_list(['ATTACK.TYPE', args.attack_type]) 107 | if args.obj_type is not None: 108 | cfg_from_list(['ATTACK.OBJ_TYPE', args.obj_type]) 109 | if args.attack2_type is not None: 110 | cfg_from_list(['ATTACK2.TYPE', args.attack2_type]) 111 | if args.obj2_type is not None: 112 | cfg_from_list(['ATTACK2.OBJ_TYPE', args.obj2_type]) 113 | if args.warm_num_iter is not None: 114 | cfg_from_list(['ATTACK2.STEP', args.warm_num_iter]) 115 | if args.inv is not None: 116 | cfg_from_list(['ATTACK.INV', args.inv]) 117 | if args.black is not None: 118 | cfg_from_list(['ATTACK.BLACK', args.black]) 119 | if args.pretrained_path is not None: 120 | cfg_from_list(['PRETRAINED_PATH', args.pretrained_path]) 121 | 122 | if args.mode is not None: 123 | cfg_from_list(['TRAIN.MODE', args.mode]) 124 | if args.reg_level is not None: 125 | cfg_from_list(['TRAIN.REG_LEVEL', args.reg_level]) 126 | if args.reg_ratio is not None: 127 | cfg_from_list(['TRAIN.REG_RATIO', args.reg_ratio]) 128 | if args.burn_in_period is not None: 129 | cfg_from_list(['TRAIN.BURN_IN_PERIOD', args.burn_in_period]) 130 | 131 | if len(cfg.MODEL_NAME) != 0 and len(cfg.DATASET_NAME) != 0: 132 | # outp_path = get_output_dir(cfg.MODEL_NAME, cfg.DATASET_NAME) 133 | inv_str = 'Inv-' if args.inv else '' 134 | black_str = 'BB-' if args.black else '' 135 | if cfg.TRAIN.MODE == 'at': 136 | mode_str = cfg.TRAIN.MODE 137 | attack_str = black_str+inv_str+cfg.ATTACK.LOSS_FUNC+'_'+cfg.ATTACK.OBJ_TYPE+'_'+cfg.ATTACK.TYPE 138 | elif cfg.TRAIN.MODE == '2step': 139 | mode_str = cfg.TRAIN.MODE+'_warmup_'+str(cfg.TRAIN.BURN_IN_PERIOD)+'_'+str(cfg.TRAIN.REG_RATIO) 140 | attack_str = black_str+inv_str+cfg.ATTACK.LOSS_FUNC+'_att1_'+cfg.ATTACK.OBJ_TYPE+'_'+cfg.ATTACK.TYPE+'_s_'+str(cfg.ATTACK.STEP)+'_att2_'+cfg.ATTACK2.OBJ_TYPE+'_'+cfg.ATTACK2.TYPE+'_s_'+str(cfg.ATTACK2.STEP) 141 | elif cfg.TRAIN.MODE == 'eval': 142 | mode_str = cfg.TRAIN.MODE+'_'+cfg.EVAL.MODE 143 | attack_str = black_str+inv_str+cfg.ATTACK.LOSS_FUNC+'_'+cfg.ATTACK.OBJ_TYPE+'_'+cfg.ATTACK.TYPE+'_s_'+str(cfg.ATTACK.EVAL_STEP) 144 | 145 | prefix_head=args.prefix+'_'+mode_str+'_'+attack_str+'_'+cfg.MODEL_NAME+'_'+cfg.DATASET_NAME 146 | 147 | outp_path = get_output_dir_new(prefix_head, now_day=args.datetime, exp_name=args.exp_name) 148 | cfg_from_list(['OUTPUT_PATH', outp_path]) 149 | #assert len(cfg.OUTPUT_PATH) != 0, 'Invalid OUTPUT_PATH! Make sure model name and dataset name are specified.' 150 | if not Path(cfg.OUTPUT_PATH).exists(): 151 | Path(cfg.OUTPUT_PATH).mkdir(parents=True) 152 | 153 | return args 154 | -------------------------------------------------------------------------------- /models/NGM/hypermodel.py: -------------------------------------------------------------------------------- 1 | import torch.nn.functional as F 2 | 3 | from src.lap_solvers.sinkhorn import Sinkhorn 4 | from src.lap_solvers.hungarian import hungarian 5 | from src.build_graphs import reshape_edge_feature 6 | from src.feature_align import feature_align 7 | from src.factorize_graph_matching import construct_aff_mat 8 | from models.NGM.gnn import HyperGNNLayer 9 | from models.NGM.geo_edge_feature import geo_edge_feature 10 | from models.GMN.affinity_layer import GaussianAffinity, InnerpAffinity 11 | 12 | from src.utils.config import cfg 13 | 14 | from src.backbone import * 15 | CNN = eval(cfg.BACKBONE) 16 | 17 | 18 | class Net(CNN): 19 | def __init__(self): 20 | super(Net, self).__init__() 21 | self.geo_affinity_layer = GaussianAffinity(1, cfg.NGM.GAUSSIAN_SIGMA) 22 | self.feat_affinity_layer = InnerpAffinity(cfg.NGM.FEATURE_CHANNEL) 23 | self.feat_affinity_layer3 = InnerpAffinity(cfg.NGM.FEATURE_CHANNEL) 24 | self.tau = cfg.NGM.SK_TAU 25 | self.bi_stochastic = Sinkhorn(max_iter=cfg.NGM.SK_ITER_NUM, tau=self.tau, epsilon=cfg.NGM.SK_EPSILON) 26 | self.l2norm = nn.LocalResponseNorm(cfg.NGM.FEATURE_CHANNEL * 2, alpha=cfg.NGM.FEATURE_CHANNEL * 2, beta=0.5, k=0) 27 | 28 | self.gnn_layer = cfg.NGM.GNN_LAYER 29 | for i in range(self.gnn_layer): 30 | if i == 0: 31 | gnn_layer = HyperGNNLayer( 32 | 1, 1, cfg.NGM.GNN_FEAT[i] + (1 if cfg.NGM.SK_EMB else 0), cfg.NGM.GNN_FEAT[i], 33 | sk_channel=cfg.NGM.SK_EMB, sk_tau=self.tau 34 | ) 35 | else: 36 | gnn_layer = HyperGNNLayer( 37 | cfg.NGM.GNN_FEAT[i - 1] + (1 if cfg.NGM.SK_EMB else 0), cfg.NGM.GNN_FEAT[i - 1], 38 | cfg.NGM.GNN_FEAT[i] + (1 if cfg.NGM.SK_EMB else 0), cfg.NGM.GNN_FEAT[i], 39 | sk_channel=cfg.NGM.SK_EMB, sk_tau=self.tau 40 | ) 41 | self.add_module('gnn_layer_{}'.format(i), gnn_layer) 42 | 43 | self.classifier = nn.Linear(cfg.NGM.GNN_FEAT[-1] + (1 if cfg.NGM.SK_EMB else 0), 1) 44 | 45 | self.weight2 = cfg.NGM.WEIGHT2 46 | self.weight3 = cfg.NGM.WEIGHT3 47 | self.rescale = cfg.PROBLEM.RESCALE 48 | 49 | def forward(self, data_dict): 50 | if 'images' in data_dict: 51 | # real image data 52 | src, tgt = data_dict['images'] 53 | P_src, P_tgt = data_dict['Ps'] 54 | ns_src, ns_tgt = data_dict['ns'] 55 | G_src, G_tgt = data_dict['Gs'] 56 | H_src, H_tgt = data_dict['Hs'] 57 | K_G, K_H = data_dict['KGHs'] 58 | 59 | # extract feature 60 | src_node = self.node_layers(src) 61 | src_edge = self.edge_layers(src_node) 62 | tgt_node = self.node_layers(tgt) 63 | tgt_edge = self.edge_layers(tgt_node) 64 | 65 | # feature normalization 66 | src_node = self.l2norm(src_node) 67 | src_edge = self.l2norm(src_edge) 68 | tgt_node = self.l2norm(tgt_node) 69 | tgt_edge = self.l2norm(tgt_edge) 70 | 71 | # arrange features 72 | U_src = feature_align(src_node, P_src, ns_src, self.rescale) 73 | F_src = feature_align(src_edge, P_src, ns_src, self.rescale) 74 | U_tgt = feature_align(tgt_node, P_tgt, ns_tgt, self.rescale) 75 | F_tgt = feature_align(tgt_edge, P_tgt, ns_tgt, self.rescale) 76 | elif 'features' in data_dict: 77 | # synthetic data 78 | src, tgt = data_dict['features'] 79 | P_src, P_tgt = data_dict['Ps'] 80 | ns_src, ns_tgt = data_dict['ns'] 81 | G_src, G_tgt = data_dict['Gs'] 82 | H_src, H_tgt = data_dict['Hs'] 83 | K_G, K_H = data_dict['KGHs'] 84 | 85 | U_src = src[:, :src.shape[1] // 2, :] 86 | F_src = src[:, src.shape[1] // 2:, :] 87 | U_tgt = tgt[:, :tgt.shape[1] // 2, :] 88 | F_tgt = tgt[:, tgt.shape[1] // 2:, :] 89 | else: 90 | raise ValueError('unknown type string {}'.format(type)) 91 | 92 | X = reshape_edge_feature(F_src, G_src, H_src) 93 | Y = reshape_edge_feature(F_tgt, G_tgt, H_tgt) 94 | dx = geo_edge_feature(P_src, G_src, H_src)[:, :1, :] 95 | dy = geo_edge_feature(P_tgt, G_tgt, H_tgt)[:, :1, :] 96 | 97 | # affinity layer for 2-order affinity matrix 98 | if cfg.NGM.EDGE_FEATURE == 'cat': 99 | Ke, Kp = self.feat_affinity_layer(X, Y, U_src, U_tgt) 100 | elif cfg.NGM.EDGE_FEATURE == 'geo': 101 | Ke, Kp = self.geo_affinity_layer(dx, dy, U_src, U_tgt) 102 | else: 103 | raise ValueError('Unknown edge feature type {}'.format(cfg.NGM.EDGE_FEATURE)) 104 | 105 | K = construct_aff_mat(Ke, torch.zeros_like(Kp), K_G, K_H) 106 | adj = (K > 0).to(K.dtype) 107 | 108 | # build 3-order affinity tensor 109 | hshape = list(adj.shape) + [adj.shape[-1]] 110 | order3A = adj.unsqueeze(1).expand(hshape) * adj.unsqueeze(2).expand(hshape) * adj.unsqueeze(3).expand(hshape) 111 | hyper_adj = order3A 112 | 113 | if cfg.NGM.ORDER3_FEATURE == 'cat': 114 | Ke_3, _ = self.feat_affinity_layer3(X, Y, torch.zeros(1, 1, 1), torch.zeros(1, 1, 1), w1=0.5, w2=1) 115 | K_3 = construct_aff_mat(Ke_3, torch.zeros_like(Kp), K_G, K_H) 116 | H = (K_3.unsqueeze(1).expand(hshape) + K_3.unsqueeze(2).expand(hshape) + K_3.unsqueeze(3).expand(hshape)) * F.relu(self.weight3) 117 | elif cfg.NGM.ORDER3_FEATURE == 'geo': 118 | Ke_d, _ = self.geo_affinity_layer(dx, dy, torch.zeros(1, 1, 1), torch.zeros(1, 1, 1)) 119 | 120 | m_d_src = construct_aff_mat(dx.squeeze().unsqueeze(-1).expand_as(Ke_d), torch.zeros_like(Kp), K_G, K_H).cpu() 121 | m_d_tgt = construct_aff_mat(dy.squeeze().unsqueeze(-2).expand_as(Ke_d), torch.zeros_like(Kp), K_G, K_H).cpu() 122 | order3A = order3A.cpu() 123 | 124 | cum_sin = torch.zeros_like(order3A) 125 | for i in range(3): 126 | def calc_sin(t): 127 | a = t.unsqueeze(i % 3 + 1).expand(hshape) 128 | b = t.unsqueeze((i + 1) % 3 + 1).expand(hshape) 129 | c = t.unsqueeze((i + 2) % 3 + 1).expand(hshape) 130 | cos = torch.clamp((a.pow(2) + b.pow(2) - c.pow(2)) / (2 * a * b + 1e-15), -1, 1) 131 | cos *= order3A 132 | sin = torch.sqrt(1 - cos.pow(2)) * order3A 133 | assert torch.sum(torch.isnan(sin)) == 0 134 | return sin 135 | sin_src = calc_sin(m_d_src) 136 | sin_tgt = calc_sin(m_d_tgt) 137 | cum_sin += torch.abs(sin_src - sin_tgt) 138 | 139 | H = torch.exp(- 1 / cfg.NGM.SIGMA3 * cum_sin) * order3A 140 | H = H.cuda() 141 | order3A = order3A.cuda() 142 | elif cfg.NGM.ORDER3_FEATURE == 'none': 143 | H = torch.zeros_like(hyper_adj) 144 | else: 145 | raise ValueError('Unknown edge feature type {}'.format(cfg.NGM.ORDER3_FEATURE)) 146 | 147 | hyper_adj = hyper_adj.cpu() 148 | hyper_adj_sum = torch.sum(hyper_adj, dim=tuple(range(2, 3 + 1)), keepdim=True) + 1e-10 149 | hyper_adj = hyper_adj / hyper_adj_sum 150 | hyper_adj = hyper_adj.to_sparse().coalesce().cuda() 151 | 152 | H = H.sparse_mask(hyper_adj) 153 | H = (H._indices(), H._values().unsqueeze(-1)) 154 | 155 | if cfg.NGM.FIRST_ORDER: 156 | emb = Kp.transpose(1, 2).contiguous().view(Kp.shape[0], -1, 1) 157 | else: 158 | emb = torch.ones(K.shape[0], K.shape[1], 1, device=K.device) 159 | 160 | adj_sum = torch.sum(adj, dim=2, keepdim=True) + 1e-10 161 | adj = adj / adj_sum 162 | pack_M = [K.unsqueeze(-1), H] 163 | pack_A = [adj, hyper_adj] 164 | for i in range(self.gnn_layer): 165 | gnn_layer = getattr(self, 'gnn_layer_{}'.format(i)) 166 | pack_M, emb = gnn_layer(pack_A, pack_M, emb, ns_src, ns_tgt, norm=False) 167 | 168 | v = self.classifier(emb) 169 | s = v.view(v.shape[0], P_tgt.shape[1], -1).transpose(1, 2) 170 | 171 | ss = self.bi_stochastic(s, ns_src, ns_tgt) 172 | x = hungarian(ss, ns_src, ns_tgt) 173 | 174 | data_dict.update({ 175 | 'ds_mat': ss, 176 | 'perm_mat': x, 177 | 'aff_mat': K 178 | }) 179 | 180 | return data_dict 181 | -------------------------------------------------------------------------------- /src/build_graphs.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import Tensor 3 | from scipy.spatial import Delaunay 4 | from scipy.spatial.qhull import QhullError 5 | import scipy as sp 6 | import scipy.spatial as sptl 7 | import scipy.sparse as sprs 8 | 9 | import itertools 10 | import numpy as np 11 | 12 | from typing import Tuple 13 | 14 | 15 | def build_graphs(P_np: np.ndarray, n: int, n_pad: int=None, edge_pad: int=None, stg: str='fc', sym: bool=True, 16 | thre: int=0, namelist=None) -> Tuple[np.ndarray, np.ndarray, np.ndarray, int]: 17 | r""" 18 | Build graph matrix :math:`\mathbf G, \mathbf H` from point set :math:`\mathbf P`. 19 | This function supports only cpu operations in numpy. 20 | :math:`\mathbf G, \mathbf H` are constructed from adjacency matrix :math:`\mathbf A`: 21 | :math:`\mathbf A = \mathbf G \cdot \mathbf H^\top` 22 | 23 | :param P_np: :math:`(n\times 2)` point set containing point coordinates 24 | :param n: number of exact points in the point set 25 | :param n_pad: padded node length 26 | :param edge_pad: padded edge length 27 | :param stg: strategy to build graphs. Options: ``fc``, ``near``, ``tri`` 28 | :param sym: True for a symmetric adjacency, False for half adjacency (A contains only the upper half) 29 | :param thre: The threshold value of 'near' strategy 30 | :return: :math:`A`, :math:`G`, :math:`H`, edge_num 31 | 32 | The possible options for ``stg``: 33 | :: 34 | 35 | 'fc'(default): construct a fully-connected graph 36 | 'near': construct a fully-connected graph, but edges longer than ``thre`` are removed 37 | 'tri': apply Delaunay triangulation 38 | 39 | An illustration of :math:`\mathbf G, \mathbf H` with their connections to the graph, the adjacency matrix, 40 | the incident matrix is 41 | 42 | .. image:: ../../images/build_graphs_GH.png 43 | """ 44 | 45 | assert stg in ('fc', 'tri', 'near', 'gabriel'), 'No strategy named {} found.'.format(stg) 46 | 47 | if stg == 'tri': 48 | A = delaunay_triangulate(P_np[0:n, :]) 49 | elif stg =='gabriel': 50 | A = gabriel_graph(P_np[0:n, :]) 51 | elif stg == 'near': 52 | A = fully_connect(P_np[0:n, :], thre=thre) 53 | else: 54 | A = fully_connect(P_np[0:n, :]) 55 | edge_num = int(np.sum(A, axis=(0, 1))) 56 | assert n > 0 and edge_num > 0, 'Error in n = {} and edge_num = {}'.format(n, edge_num) 57 | 58 | if n_pad is None: 59 | n_pad = n 60 | if edge_pad is None: 61 | edge_pad = edge_num 62 | assert n_pad >= n 63 | assert edge_pad >= edge_num 64 | 65 | G = np.zeros((n_pad, edge_pad), dtype=np.float32) 66 | H = np.zeros((n_pad, edge_pad), dtype=np.float32) 67 | edge_idx = 0 68 | for i in range(n): 69 | if sym: 70 | range_j = range(n) 71 | else: 72 | range_j = range(i, n) 73 | for j in range_j: 74 | if A[i, j] == 1: 75 | G[i, edge_idx] = 1 76 | H[j, edge_idx] = 1 77 | edge_idx += 1 78 | 79 | return A, G, H, edge_num 80 | 81 | 82 | def delaunay_triangulate(P: np.ndarray) -> np.ndarray: 83 | r""" 84 | Perform delaunay triangulation on point set P. 85 | 86 | :param P: :math:`(n\times 2)` point set 87 | :return: adjacency matrix :math:`A` 88 | """ 89 | n = P.shape[0] 90 | if n < 3: 91 | A = fully_connect(P) 92 | else: 93 | try: 94 | d = Delaunay(P) 95 | #assert d.coplanar.size == 0, 'Delaunay triangulation omits points.' 96 | A = np.zeros((n, n)) 97 | for simplex in d.simplices: 98 | for pair in itertools.permutations(simplex, 2): 99 | A[pair] = 1 100 | except QhullError as err: 101 | print('Delaunay triangulation error detected. Return fully-connected graph.') 102 | print('Traceback:') 103 | print(err) 104 | A = fully_connect(P) 105 | return A 106 | 107 | def gabriel_graph(P: np.ndarray) -> np.ndarray: 108 | r""" 109 | Perform delaunay triangulation on point set P. 110 | 111 | :param P: :math:`(n\times 2)` point set 112 | :return: adjacency matrix :math:`A` 113 | """ 114 | n = P.shape[0] 115 | if n < 3: 116 | A = fully_connect(P) 117 | else: 118 | try: 119 | d = Delaunay(P) 120 | #assert d.coplanar.size == 0, 'Delaunay triangulation omits points.' 121 | A = np.zeros((n, n)) 122 | for simplex in d.simplices: 123 | for pair in itertools.permutations(simplex, 2): 124 | A[pair] = 1 125 | 126 | # pruning the non-Gabriel edges 127 | edge_pair_idx = np.stack(A.nonzero()).T 128 | center = d.points[edge_pair_idx] 129 | m = (center[:, 0, :] + center[:, 1, :])/2 130 | # Find the radius sphere between each pair of nodes 131 | r = sp.sqrt(sp.sum((center[:, 0, :] - center[:, 1, :])**2, axis=1))/2 132 | 133 | # Use the kd-tree function in Scipy's spatial module 134 | tree = sptl.cKDTree(P) 135 | # Find the nearest point for each midpoint 136 | n = tree.query(x=m, k=1)[0] 137 | # If nearest point to m is at a distance r, then the edge is a Gabriel edge 138 | g = n <= r*(0.999) # The factor is to avoid precision errors in the distances 139 | # Reduce the connectivity to all True values found in g 140 | A[edge_pair_idx[g][:,0], edge_pair_idx[g][:, 1]] = 0 141 | 142 | except QhullError as err: 143 | print('Delaunay triangulation error detected. Return fully-connected graph.') 144 | print('Traceback:') 145 | print(err) 146 | A = fully_connect(P) 147 | 148 | return A 149 | 150 | 151 | def fully_connect(P: np.ndarray, thre=None) -> np.ndarray: 152 | r""" 153 | Return the adjacency matrix of a fully-connected graph. 154 | 155 | :param P: :math:`(n\times 2)` point set 156 | :param thre: edges that are longer than this threshold will be removed 157 | :return: adjacency matrix :math:`A` 158 | """ 159 | n = P.shape[0] 160 | A = np.ones((n, n)) - np.eye(n) 161 | if thre is not None: 162 | for i in range(n): 163 | for j in range(i): 164 | if np.linalg.norm(P[i] - P[j]) > thre: 165 | A[i, j] = 0 166 | A[j, i] = 0 167 | return A 168 | 169 | def make_grids(start, stop, num) -> np.ndarray: 170 | r""" 171 | Make grids. 172 | 173 | This function supports only cpu operations in numpy. 174 | 175 | :param start: start index in all dimensions 176 | :param stop: stop index in all dimensions 177 | :param num: number of grids in each dimension 178 | :return: point set P 179 | """ 180 | length = np.prod(num) 181 | P = np.zeros((length, len(num)), dtype=np.float32) 182 | assert len(start) == len(stop) == len(num) 183 | for i, (begin, end, n) in enumerate(zip(start, stop, num)): 184 | g = np.linspace(begin, end, n + 1) 185 | g -= (g[1] - g[0]) / 2 186 | g = g[1:] 187 | P[:, i] = np.reshape(np.repeat([g], length / n, axis=i), length) 188 | return P 189 | 190 | 191 | def reshape_edge_feature(F: Tensor, G: Tensor, H: Tensor, device=None) -> Tensor: 192 | r""" 193 | Given point-level features extracted from images, reshape it into edge feature matrix :math:`X`, 194 | where features are arranged by the order of :math:`G`, :math:`H`. 195 | 196 | .. math:: 197 | \mathbf{X}_{e_{ij}} = concat(\mathbf{F}_i, \mathbf{F}_j) 198 | 199 | where :math:`e_{ij}` means an edge connecting nodes :math:`i, j` 200 | 201 | :param F: :math:`(b\times d \times n)` extracted point-level feature matrix. 202 | :math:`b`: batch size. :math:`d`: feature dimension. :math:`n`: number of nodes. 203 | :param G: :math:`(b\times n \times e)` factorized adjacency matrix, where :math:`\mathbf A = \mathbf G \cdot \mathbf H^\top`. :math:`e`: number of edges. 204 | :param H: :math:`(b\times n \times e)` factorized adjacency matrix, where :math:`\mathbf A = \mathbf G \cdot \mathbf H^\top` 205 | :param device: device. If not specified, it will be the same as the input 206 | :return: edge feature matrix X :math:`(b \times 2d \times e)` 207 | """ 208 | if device is None: 209 | device = F.device 210 | 211 | batch_num = F.shape[0] 212 | feat_dim = F.shape[1] 213 | point_num, edge_num = G.shape[1:3] 214 | X = torch.zeros(batch_num, 2 * feat_dim, edge_num, dtype=torch.float32, device=device) 215 | X[:, 0:feat_dim, :] = torch.matmul(F, G) 216 | X[:, feat_dim:2*feat_dim, :] = torch.matmul(F, H) 217 | 218 | return X 219 | -------------------------------------------------------------------------------- /src/spectral_clustering.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | from torch import Tensor 4 | from typing import Union, Tuple 5 | 6 | def initialize(X: Tensor, num_clusters: int, method: str='plus') -> np.array: 7 | r""" 8 | Initialize cluster centers. 9 | 10 | :param X: matrix 11 | :param num_clusters: number of clusters 12 | :param method: denotes different initialization strategies: ``'plus'`` (default) or ``'random'`` 13 | :return: initial state 14 | 15 | .. note:: 16 | We support two initialization strategies: random initialization by setting ``method='random'``, or `kmeans++ 17 | `_ by setting ``method='plus'``. 18 | """ 19 | if method == 'plus': 20 | init_func = _initialize_plus 21 | elif method == 'random': 22 | init_func = _initialize_random 23 | else: 24 | raise NotImplementedError 25 | return init_func(X, num_clusters) 26 | 27 | 28 | def _initialize_random(X, num_clusters): 29 | """ 30 | Initialize cluster centers randomly. See :func:`src.spectral_clustering.initialize` for details. 31 | """ 32 | num_samples = len(X) 33 | indices = np.random.choice(num_samples, num_clusters, replace=False) 34 | initial_state = X[indices] 35 | return initial_state 36 | 37 | def _initialize_plus(X, num_clusters): 38 | """ 39 | Initialize cluster centers by k-means++. See :func:`src.spectral_clustering.initialize` for details. 40 | """ 41 | num_samples = len(X) 42 | centroid_index = np.zeros(num_clusters) 43 | for i in range(num_clusters): 44 | if i == 0: 45 | choice_prob = np.full(num_samples, 1 / num_samples) 46 | else: 47 | centroid_X = X[centroid_index[:i]] 48 | dis = _pairwise_distance(X, centroid_X) 49 | dis_to_nearest_centroid = torch.min(dis, dim=1).values 50 | choice_prob = dis_to_nearest_centroid / torch.sum(dis_to_nearest_centroid) 51 | choice_prob = choice_prob.detach().cpu().numpy() 52 | 53 | centroid_index[i] = np.random.choice(num_samples, 1, p=choice_prob, replace=False) 54 | 55 | initial_state = X[centroid_index] 56 | return initial_state 57 | 58 | def kmeans( 59 | X: Tensor, 60 | num_clusters: int, 61 | init_x: Union[Tensor, str]='plus', 62 | distance: str='euclidean', 63 | tol: float=1e-4, 64 | device=torch.device('cpu'), 65 | ) -> Tuple[Tensor, Tensor]: 66 | r""" 67 | Perform kmeans on given data matrix :math:`\mathbf X`. 68 | 69 | :param X: :math:`(n\times d)` input data matrix. :math:`n`: number of samples. :math:`d`: feature dimension 70 | :param num_clusters: (int) number of clusters 71 | :param init_x: how to initiate x (provide a initial state of x or define a init method) [default: 'plus'] 72 | :param distance: distance [options: 'euclidean', 'cosine'] [default: 'euclidean'] 73 | :param tol: convergence threshold [default: 0.0001] 74 | :param device: computing device [default: cpu] 75 | :return: cluster ids, cluster centers 76 | """ 77 | if distance == 'euclidean': 78 | pairwise_distance_function = _pairwise_distance 79 | elif distance == 'cosine': 80 | pairwise_distance_function = _pairwise_cosine 81 | else: 82 | raise NotImplementedError 83 | 84 | # convert to float 85 | X = X.float() 86 | 87 | # transfer to device 88 | X = X.to(device) 89 | 90 | # initialize 91 | if type(init_x) is str: 92 | initial_state = initialize(X, num_clusters, method=init_x) 93 | else: 94 | initial_state = init_x 95 | 96 | iteration = 0 97 | while True: 98 | dis = pairwise_distance_function(X, initial_state) 99 | 100 | choice_cluster = torch.argmin(dis, dim=1) 101 | 102 | initial_state_pre = initial_state.clone() 103 | 104 | for index in range(num_clusters): 105 | selected = torch.nonzero(choice_cluster == index, as_tuple=False).squeeze().to(device) 106 | 107 | selected = torch.index_select(X, 0, selected) 108 | initial_state[index] = selected.mean(dim=0) 109 | 110 | center_shift = torch.sum( 111 | torch.sqrt( 112 | torch.sum((initial_state - initial_state_pre) ** 2, dim=1) 113 | )) 114 | 115 | # increment iteration 116 | iteration = iteration + 1 117 | 118 | if center_shift ** 2 < tol: 119 | break 120 | 121 | if torch.isnan(initial_state).any(): 122 | print('NAN encountered in clustering. Retrying...') 123 | initial_state = initialize(X, num_clusters) 124 | 125 | return choice_cluster.cpu(), initial_state.cpu() 126 | 127 | 128 | def kmeans_predict( 129 | X: Tensor, 130 | cluster_centers: Tensor, 131 | distance: str='euclidean', 132 | device=torch.device('cpu') 133 | ) -> Tensor: 134 | r""" 135 | Kmeans prediction using existing cluster centers. 136 | 137 | :param X: matrix 138 | :param cluster_centers: cluster centers 139 | :param distance: distance [options: 'euclidean', 'cosine'] [default: 'euclidean'] 140 | :param device: computing device [default: 'cpu'] 141 | :return: cluster ids 142 | """ 143 | if distance == 'euclidean': 144 | pairwise_distance_function = _pairwise_distance 145 | elif distance == 'cosine': 146 | pairwise_distance_function = _pairwise_cosine 147 | else: 148 | raise NotImplementedError 149 | 150 | # convert to float 151 | X = X.float() 152 | 153 | # transfer to device 154 | X = X.to(device) 155 | 156 | dis = pairwise_distance_function(X, cluster_centers) 157 | choice_cluster = torch.argmin(dis, dim=1) 158 | 159 | return choice_cluster.cpu() 160 | 161 | 162 | def _pairwise_distance(data1, data2, device=torch.device('cpu')): 163 | """Compute pairwise Euclidean distance""" 164 | # transfer to device 165 | data1, data2 = data1.to(device), data2.to(device) 166 | 167 | # N*1*M 168 | A = data1.unsqueeze(dim=1) 169 | 170 | # 1*N*M 171 | B = data2.unsqueeze(dim=0) 172 | 173 | dis = (A - B) ** 2.0 174 | # return N*N matrix for pairwise distance 175 | dis = dis.sum(dim=-1) #.squeeze(-1) 176 | return dis 177 | 178 | 179 | def _pairwise_cosine(data1, data2, device=torch.device('cpu')): 180 | """Compute pairwise cosine distance""" 181 | # transfer to device 182 | data1, data2 = data1.to(device), data2.to(device) 183 | 184 | # N*1*M 185 | A = data1.unsqueeze(dim=1) 186 | 187 | # 1*N*M 188 | B = data2.unsqueeze(dim=0) 189 | 190 | # normalize the points | [0.3, 0.4] -> [0.3/sqrt(0.09 + 0.16), 0.4/sqrt(0.09 + 0.16)] = [0.3/0.5, 0.4/0.5] 191 | A_normalized = A / A.norm(dim=-1, keepdim=True) 192 | B_normalized = B / B.norm(dim=-1, keepdim=True) 193 | 194 | cosine = A_normalized * B_normalized 195 | 196 | # return N*N matrix for pairwise distance 197 | cosine_dis = 1 - cosine.sum(dim=-1).squeeze(-1) 198 | return cosine_dis 199 | 200 | 201 | def spectral_clustering(sim_matrix: Tensor, cluster_num: int, init: Tensor=None, 202 | return_state: bool=False, normalized: bool=False): 203 | r""" 204 | Perform spectral clustering based on given similarity matrix. 205 | 206 | This function firstly computes the leading eigenvectors of the given similarity matrix, and then utilizes the 207 | eigenvectors as features and performs k-means clustering based on these features. 208 | 209 | :param sim_matrix: :math:`(n\times n)` input similarity matrix. :math:`n`: number of instances 210 | :param cluster_num: number of clusters 211 | :param init: the initialization technique or initial features for k-means 212 | :param return_state: whether return state features (can be further used for prediction) 213 | :param normalized: whether to normalize the similarity matrix by its degree 214 | :return: the belonging of each instance to clusters, state features (if ``return_state==True``) 215 | """ 216 | degree = torch.diagflat(torch.sum(sim_matrix, dim=-1)) 217 | if normalized: 218 | aff_matrix = (degree - sim_matrix) / torch.diag(degree).unsqueeze(1) 219 | else: 220 | aff_matrix = degree - sim_matrix 221 | e, v = torch.symeig(aff_matrix, eigenvectors=True) 222 | topargs = torch.argsort(torch.abs(e), descending=False)[1:cluster_num] 223 | v = v[:, topargs] 224 | 225 | if cluster_num == 2: 226 | choice_cluster = (v > 0).to(torch.int).squeeze(1) 227 | else: 228 | choice_cluster, initial_state = kmeans(v, cluster_num, init if init is not None else 'plus', 229 | distance='euclidean', tol=1e-6) 230 | 231 | choice_cluster = choice_cluster.to(sim_matrix.device) 232 | 233 | if return_state: 234 | return choice_cluster, initial_state 235 | else: 236 | return choice_cluster 237 | 238 | -------------------------------------------------------------------------------- /src/dataset/willow_obj.py: -------------------------------------------------------------------------------- 1 | from pathlib import Path 2 | import scipy.io as sio 3 | from PIL import Image 4 | import numpy as np 5 | from src.utils.config import cfg 6 | from src.dataset.base_dataset import BaseDataset 7 | import random 8 | 9 | 10 | ''' 11 | Important Notice: Face image 160 contains only 8 labeled keypoints (should be 10) 12 | ''' 13 | 14 | class WillowObject(BaseDataset): 15 | def __init__(self, sets, obj_resize): 16 | """ 17 | :param sets: 'train' or 'test' 18 | :param obj_resize: resized object size 19 | """ 20 | super(WillowObject, self).__init__() 21 | self.classes = cfg.WillowObject.CLASSES 22 | self.kpt_len = [cfg.WillowObject.KPT_LEN for _ in cfg.WillowObject.CLASSES] 23 | 24 | self.root_path = Path(cfg.WillowObject.ROOT_DIR) 25 | self.obj_resize = obj_resize 26 | 27 | assert sets in ('train', 'test'), 'No match found for dataset {}'.format(sets) 28 | self.sets = sets 29 | self.split_offset = cfg.WillowObject.SPLIT_OFFSET 30 | self.train_len = cfg.WillowObject.TRAIN_NUM 31 | self.rand_outlier = cfg.WillowObject.RAND_OUTLIER 32 | 33 | self.mat_list = [] 34 | for cls_name in self.classes: 35 | assert type(cls_name) is str 36 | cls_mat_list = [p for p in (self.root_path / cls_name).glob('*.mat')] 37 | if cls_name == 'Face': 38 | cls_mat_list.remove(self.root_path / cls_name / 'image_0160.mat') 39 | assert not self.root_path / cls_name / 'image_0160.mat' in cls_mat_list 40 | ori_len = len(cls_mat_list) 41 | if self.split_offset % ori_len + self.train_len <= ori_len: 42 | if sets == 'train' and not cfg.WillowObject.TRAIN_SAME_AS_TEST: 43 | self.mat_list.append( 44 | cls_mat_list[self.split_offset % ori_len: (self.split_offset + self.train_len) % ori_len] 45 | ) 46 | else: 47 | self.mat_list.append( 48 | cls_mat_list[:self.split_offset % ori_len] + 49 | cls_mat_list[(self.split_offset + self.train_len) % ori_len:] 50 | ) 51 | else: 52 | if sets == 'train' and not cfg.WillowObject.TRAIN_SAME_AS_TEST: 53 | self.mat_list.append( 54 | cls_mat_list[:(self.split_offset + self.train_len) % ori_len - ori_len] + 55 | cls_mat_list[self.split_offset % ori_len:] 56 | ) 57 | else: 58 | self.mat_list.append( 59 | cls_mat_list[(self.split_offset + self.train_len) % ori_len - ori_len: self.split_offset % ori_len] 60 | ) 61 | 62 | def get_pair(self, cls=None, shuffle=True): 63 | """ 64 | Randomly get a pair of objects from WILLOW-object dataset 65 | :param cls: None for random class, or specify for a certain set 66 | :param shuffle: random shuffle the keypoints 67 | :return: (pair of data, groundtruth permutation matrix) 68 | """ 69 | if cls is None: 70 | cls = random.randrange(0, len(self.classes)) 71 | elif type(cls) == str: 72 | cls = self.classes.index(cls) 73 | assert type(cls) == int and 0 <= cls < len(self.classes) 74 | 75 | anno_pair = [] 76 | for mat_name in random.sample(self.mat_list[cls], 2): 77 | anno_dict = self.__get_anno_dict(mat_name, cls) 78 | if shuffle: 79 | random.shuffle(anno_dict['keypoints']) 80 | anno_pair.append(anno_dict) 81 | 82 | perm_mat = np.zeros([len(_['keypoints']) for _ in anno_pair], dtype=np.float32) 83 | row_list = [] 84 | col_list = [] 85 | for i, keypoint in enumerate(anno_pair[0]['keypoints']): 86 | for j, _keypoint in enumerate(anno_pair[1]['keypoints']): 87 | if keypoint['name'] == _keypoint['name']: 88 | if keypoint['name'] != 'outlier': 89 | perm_mat[i, j] = 1 90 | row_list.append(i) 91 | col_list.append(j) 92 | break 93 | row_list.sort() 94 | col_list.sort() 95 | perm_mat = perm_mat[row_list, :] 96 | perm_mat = perm_mat[:, col_list] 97 | anno_pair[0]['keypoints'] = [anno_pair[0]['keypoints'][i] for i in row_list] 98 | anno_pair[1]['keypoints'] = [anno_pair[1]['keypoints'][j] for j in col_list] 99 | 100 | return anno_pair, perm_mat 101 | 102 | def get_multi(self, cls=None, num=2, shuffle=True): 103 | """ 104 | Randomly get multiple objects from Willow Object Class dataset for multi-matching. 105 | :param cls: None for random class, or specify for a certain set 106 | :param num: number of objects to be fetched 107 | :param shuffle: random shuffle the keypoints 108 | :return: (list of data, list of permutation matrices) 109 | """ 110 | if cls is None: 111 | cls = random.randrange(0, len(self.classes)) 112 | elif type(cls) == str: 113 | cls = self.classes.index(cls) 114 | assert type(cls) == int and 0 <= cls < len(self.classes) 115 | 116 | anno_list = [] 117 | for mat_name in random.sample(self.mat_list[cls], num): 118 | anno_dict = self.__get_anno_dict(mat_name, cls) 119 | if shuffle: 120 | random.shuffle(anno_dict['keypoints']) 121 | anno_list.append(anno_dict) 122 | 123 | perm_mat = [np.zeros([len(anno_list[0]['keypoints']), len(x['keypoints'])], dtype=np.float32) for x in 124 | anno_list] 125 | row_list = [] 126 | col_lists = [] 127 | for i in range(num): 128 | col_lists.append([]) 129 | for i, keypoint in enumerate(anno_list[0]['keypoints']): 130 | kpt_idx = [] 131 | for anno_dict in anno_list: 132 | kpt_name_list = [x['name'] for x in anno_dict['keypoints']] 133 | if keypoint['name'] in kpt_name_list: 134 | kpt_idx.append(kpt_name_list.index(keypoint['name'])) 135 | else: 136 | kpt_idx.append(-1) 137 | row_list.append(i) 138 | for k in range(num): 139 | j = kpt_idx[k] 140 | if j != -1: 141 | col_lists[k].append(j) 142 | if keypoint['name'] != 'outlier': 143 | perm_mat[k][i, j] = 1 144 | 145 | row_list.sort() 146 | for col_list in col_lists: 147 | col_list.sort() 148 | 149 | for k in range(num): 150 | perm_mat[k] = perm_mat[k][row_list, :] 151 | perm_mat[k] = perm_mat[k][:, col_lists[k]] 152 | anno_list[k]['keypoints'] = [anno_list[k]['keypoints'][j] for j in col_lists[k]] 153 | perm_mat[k] = perm_mat[k].transpose() 154 | 155 | return anno_list, perm_mat 156 | 157 | def __get_anno_dict(self, mat_file, cls): 158 | """ 159 | Get an annotation dict from .mat annotation 160 | """ 161 | assert mat_file.exists(), '{} does not exist.'.format(mat_file) 162 | 163 | img_name = mat_file.stem + '.png' 164 | img_file = mat_file.parent / img_name 165 | 166 | struct = sio.loadmat(mat_file.open('rb')) 167 | kpts = struct['pts_coord'] 168 | 169 | with Image.open(str(img_file)) as img: 170 | ori_sizes = img.size 171 | obj = img.resize(self.obj_resize, resample=Image.BICUBIC) 172 | xmin = 0 173 | ymin = 0 174 | w = ori_sizes[0] 175 | h = ori_sizes[1] 176 | 177 | keypoint_list = [] 178 | for idx, keypoint in enumerate(np.split(kpts, kpts.shape[1], axis=1)): 179 | attr = { 180 | 'name': idx, 181 | 'x': float(keypoint[0]) * self.obj_resize[0] / w, 182 | 'y': float(keypoint[1]) * self.obj_resize[1] / h 183 | } 184 | keypoint_list.append(attr) 185 | 186 | for idx in range(self.rand_outlier): 187 | attr = { 188 | 'name': 'outlier', 189 | 'x': random.uniform(0, self.obj_resize[0]), 190 | 'y': random.uniform(0, self.obj_resize[1]) 191 | } 192 | keypoint_list.append(attr) 193 | 194 | anno_dict = dict() 195 | anno_dict['image'] = obj 196 | anno_dict['keypoints'] = keypoint_list 197 | anno_dict['bounds'] = xmin, ymin, w, h 198 | anno_dict['ori_sizes'] = ori_sizes 199 | anno_dict['cls'] = cls 200 | anno_dict['univ_size'] = 10 201 | 202 | return anno_dict 203 | 204 | def len(self, cls): 205 | if type(cls) == int: 206 | cls = self.classes[cls] 207 | assert cls in self.classes 208 | return len(self.mat_list[self.classes.index(cls)]) 209 | 210 | 211 | if __name__ == '__main__': 212 | cfg.WillowObject.ROOT_DIR = 'WILLOW-ObjectClass' 213 | cfg.WillowObject.SPLIT_OFFSET = 0 214 | train = WillowObject('train', (256, 256)) 215 | test = WillowObject('test', (256, 256)) 216 | for train_cls_list, test_cls_list in zip(train.mat_list, test.mat_list): 217 | for t in train_cls_list: 218 | assert t not in test_cls_list 219 | pass 220 | -------------------------------------------------------------------------------- /src/extension/sparse_dot/sparse_dot.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | /* CUDA Declaration */ 5 | 6 | at::Tensor csr_dot_csc_cuda( 7 | at::Tensor t1_indices, 8 | at::Tensor t1_indptr, 9 | at::Tensor t1_data, 10 | at::Tensor t2_indices, 11 | at::Tensor t2_indptr, 12 | at::Tensor t2_data, 13 | int64_t batch_size, 14 | int64_t out_h, 15 | int64_t out_w 16 | ); 17 | 18 | 19 | std::vector csr_dot_diag_cuda( 20 | at::Tensor t1_indices, 21 | at::Tensor t1_indptr, 22 | at::Tensor t1_data, 23 | at::Tensor t2, 24 | int64_t batch_size, 25 | int64_t out_h, 26 | int64_t out_w 27 | ); 28 | 29 | 30 | #define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be a CUDA tensor") 31 | #define CHECK_CPU(x) AT_ASSERTM(!x.type().is_cuda(), #x " must be a CPU tensor") 32 | #define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous") 33 | #define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) 34 | 35 | 36 | /* CSR dot CSC Implementation */ 37 | 38 | std::vector csr_dot_csc_cpu( 39 | at::Tensor t1_indices, 40 | at::Tensor t1_indptr, 41 | at::Tensor t1_data, 42 | at::Tensor t2_indices, 43 | at::Tensor t2_indptr, 44 | at::Tensor t2_data, 45 | int64_t batch_size, 46 | int64_t out_h, 47 | int64_t out_w 48 | ){ 49 | CHECK_CPU(t1_indices); 50 | CHECK_CPU(t1_indptr); 51 | CHECK_CPU(t1_data); 52 | CHECK_CPU(t2_indices); 53 | CHECK_CPU(t2_indptr); 54 | CHECK_CPU(t2_data); 55 | 56 | std::list out_indices_list[batch_size * out_h]; 57 | std::list out_data_list[batch_size * out_h]; 58 | auto out_indptr = at::zeros({batch_size * out_h + 1}, t1_indptr.type()); 59 | auto t1_indptr_acc = t1_indptr.accessor(); 60 | auto t2_indptr_acc = t2_indptr.accessor(); 61 | auto t1_indices_acc = t1_indices.accessor(); 62 | auto t2_indices_acc = t2_indices.accessor(); 63 | 64 | for (int64_t b = 0; b < batch_size; b++) 65 | { 66 | for (int64_t i = 0; i < out_h; i++) 67 | { 68 | int64_t t1_start = t1_indptr_acc[b * out_h + i]; 69 | int64_t t1_stop = t1_indptr_acc[b * out_h + i + 1]; 70 | int64_t row_nnz = 0; 71 | 72 | for (int64_t j = 0; j < out_w; j++) 73 | { 74 | int64_t t2_start = t2_indptr_acc[b * out_w + j]; 75 | int64_t t2_stop = t2_indptr_acc[b * out_w + j + 1]; 76 | 77 | float outp = 0;//at::zeros({}, t1_data.type()); 78 | int64_t t1_ptr_idx = t1_start; 79 | int64_t t2_ptr_idx = t2_start; 80 | 81 | while (t1_ptr_idx < t1_stop && t2_ptr_idx < t2_stop) 82 | { 83 | int64_t t1_cur_indice = t1_indices_acc[t1_ptr_idx]; 84 | int64_t t2_cur_indice = t2_indices_acc[t2_ptr_idx]; 85 | if (t1_cur_indice == t2_cur_indice) 86 | { 87 | auto tmp = t1_data[t1_ptr_idx] * t2_data[t2_ptr_idx]; 88 | auto tmp_acc = tmp.accessor(); 89 | outp += tmp_acc[0]; 90 | t1_ptr_idx++; 91 | t2_ptr_idx++; 92 | } 93 | else if (t1_cur_indice < t2_cur_indice) 94 | t1_ptr_idx++; 95 | else 96 | t2_ptr_idx++; 97 | } 98 | if (outp != 0) 99 | { 100 | out_data_list[b * out_h + i].push_back(outp); 101 | out_indices_list[b * out_h + i].push_back(j); 102 | row_nnz++; 103 | } 104 | } 105 | out_indptr[b * out_h + i + 1] = out_indptr[b * out_h + i] + row_nnz; 106 | } 107 | } 108 | 109 | auto out_indptr_acc = out_indptr.accessor(); 110 | int64_t nnz = out_indptr_acc[-1]; 111 | auto out_indices = at::zeros({nnz}, t1_indices.type()); 112 | auto out_data = at::zeros({nnz}, t1_data.type()); 113 | int64_t idx = 0; 114 | for (int64_t b = 0; b < batch_size; b++) 115 | { 116 | for (int64_t i = 0; i < out_h; i++) 117 | { 118 | auto * tmp_indices_list = &out_indices_list[b * out_h + i]; 119 | auto * tmp_data_list = &out_data_list[b * out_h + i]; 120 | while (!tmp_indices_list->empty() && !tmp_data_list->empty()) 121 | { 122 | out_indices[idx] = tmp_indices_list->front(); 123 | tmp_indices_list->pop_front(); 124 | out_data[idx] = tmp_data_list->front(); 125 | tmp_data_list->pop_front(); 126 | idx++; 127 | } 128 | } 129 | } 130 | 131 | return {out_indices, out_indptr, out_data}; 132 | } 133 | 134 | 135 | at::Tensor csr_dot_csc_dense_cuda_wrapper( 136 | at::Tensor t1_indices, 137 | at::Tensor t1_indptr, 138 | at::Tensor t1_data, 139 | at::Tensor t2_indices, 140 | at::Tensor t2_indptr, 141 | at::Tensor t2_data, 142 | int64_t batch_size, 143 | int64_t out_h, 144 | int64_t out_w 145 | ){ 146 | CHECK_INPUT(t1_indices); 147 | CHECK_INPUT(t1_indptr); 148 | CHECK_INPUT(t1_data); 149 | CHECK_INPUT(t2_indices); 150 | CHECK_INPUT(t2_indptr); 151 | CHECK_INPUT(t2_data); 152 | return csr_dot_csc_cuda(t1_indices, t1_indptr, t1_data, 153 | t2_indices, t2_indptr, t2_data, 154 | batch_size, out_h, out_w); 155 | } 156 | 157 | 158 | std::vector csr_dot_csc( 159 | at::Tensor t1_indices, 160 | at::Tensor t1_indptr, 161 | at::Tensor t1_data, 162 | at::Tensor t2_indices, 163 | at::Tensor t2_indptr, 164 | at::Tensor t2_data, 165 | int64_t batch_size, 166 | int64_t out_h, 167 | int64_t out_w 168 | ) 169 | { 170 | if (t1_indices.type().is_cuda()) 171 | throw std::runtime_error("Unexpected cuda tensor in sparse dot sparse -> sparse computation."); 172 | else 173 | return csr_dot_csc_cpu(t1_indices, t1_indptr, t1_data, t2_indices, t2_indptr, t2_data, batch_size, out_h, out_w); 174 | } 175 | 176 | at::Tensor csr_dot_csc_dense_cuda( 177 | at::Tensor t1_indices, 178 | at::Tensor t1_indptr, 179 | at::Tensor t1_data, 180 | at::Tensor t2_indices, 181 | at::Tensor t2_indptr, 182 | at::Tensor t2_data, 183 | int64_t batch_size, 184 | int64_t out_h, 185 | int64_t out_w 186 | ) 187 | { 188 | return csr_dot_csc_dense_cuda_wrapper(t1_indices, t1_indptr, t1_data, t2_indices, t2_indptr, t2_data, 189 | batch_size, out_h, out_w); 190 | } 191 | 192 | 193 | /* CSR dot diag implementation */ 194 | 195 | std::vector csr_dot_diag_cpu( 196 | at::Tensor t1_indices, 197 | at::Tensor t1_indptr, 198 | at::Tensor t1_data, 199 | at::Tensor t2, 200 | int64_t batch_size, 201 | int64_t out_h, 202 | int64_t out_w 203 | ) 204 | { 205 | CHECK_CPU(t1_indices); 206 | CHECK_CPU(t1_indptr); 207 | CHECK_CPU(t1_data); 208 | CHECK_CPU(t2); 209 | auto outp_indices = at::clone(t1_indices); 210 | auto outp_indptr = at::clone(t1_indptr); 211 | auto outp_data = at::zeros_like(t1_data); 212 | 213 | auto t1_indptr_acc = t1_indptr.accessor(); 214 | auto t1_indices_acc = t1_indices.accessor(); 215 | 216 | for (int64_t b = 0; b < batch_size; b++) 217 | { 218 | for (int64_t i = 0; i < out_h; i++) 219 | { 220 | int64_t start = t1_indptr_acc[b * out_h + i]; 221 | int64_t stop = t1_indptr_acc[b * out_h + i + 1]; 222 | for (int64_t data_idx = start; data_idx < stop; data_idx++) 223 | { 224 | int64_t row_idx = t1_indices_acc[data_idx]; 225 | outp_data[data_idx] = t1_data[data_idx] * t2[b][row_idx]; 226 | } 227 | } 228 | } 229 | return {outp_indices, outp_indptr, outp_data}; 230 | } 231 | 232 | 233 | std::vector csr_dot_diag_cuda_wrapper( 234 | at::Tensor t1_indices, 235 | at::Tensor t1_indptr, 236 | at::Tensor t1_data, 237 | at::Tensor t2, 238 | int64_t batch_size, 239 | int64_t out_h, 240 | int64_t out_w 241 | ) 242 | { 243 | CHECK_INPUT(t1_indices); 244 | CHECK_INPUT(t1_indptr); 245 | CHECK_INPUT(t1_data); 246 | CHECK_INPUT(t2); 247 | return csr_dot_diag_cuda(t1_indices, t1_indptr, t1_data, t2, batch_size, out_h, out_w); 248 | } 249 | 250 | 251 | std::vector csr_dot_diag( 252 | at::Tensor t1_indices, 253 | at::Tensor t1_indptr, 254 | at::Tensor t1_data, 255 | at::Tensor t2, 256 | int64_t batch_size, 257 | int64_t out_h, 258 | int64_t out_w 259 | ) 260 | { 261 | if (t1_indices.type().is_cuda()) 262 | return csr_dot_diag_cuda_wrapper(t1_indices, t1_indptr, t1_data, t2, batch_size, out_h, out_w); 263 | else 264 | return csr_dot_diag_cpu(t1_indices, t1_indptr, t1_data, t2, batch_size, out_h, out_w); 265 | 266 | } 267 | 268 | /* PyBind Interface */ 269 | 270 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 271 | m.def("csr_dot_csc", &csr_dot_csc, "csr sparse matrix dot csc sparse matrix"); 272 | m.def("csr_dot_csc_dense_cuda", &csr_dot_csc_dense_cuda, 273 | "cuda implementation of csr sparse matrix dot csc sparse matrix, result is dense"); 274 | m.def("csr_dot_diag", &csr_dot_diag, "csr sparse matrix dot a diagonal of dense vector"); 275 | } 276 | -------------------------------------------------------------------------------- /models/NGM/model.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | from src.lap_solvers.sinkhorn import Sinkhorn, GumbelSinkhorn 5 | from src.build_graphs import reshape_edge_feature 6 | from src.feature_align import feature_align 7 | from src.factorize_graph_matching import construct_aff_mat 8 | from models.NGM.gnn import GNNLayer 9 | from models.NGM.geo_edge_feature import geo_edge_feature 10 | from models.GMN.affinity_layer import InnerpAffinity, GaussianAffinity 11 | from src.evaluation_metric import objective_score 12 | from src.lap_solvers.hungarian import hungarian 13 | import math 14 | from src.utils.gpu_memory import gpu_free_memory 15 | 16 | #from torch_geometric.data import Data, Batch 17 | #from torch_geometric.utils import dense_to_sparse, to_dense_batch 18 | 19 | from src.utils.config import cfg 20 | 21 | from src.backbone import * 22 | CNN = eval(cfg.BACKBONE) 23 | 24 | 25 | class Net(CNN): 26 | def __init__(self): 27 | super(Net, self).__init__() 28 | if cfg.NGM.EDGE_FEATURE == 'cat': 29 | self.affinity_layer = InnerpAffinity(cfg.NGM.FEATURE_CHANNEL) 30 | elif cfg.NGM.EDGE_FEATURE == 'geo': 31 | self.affinity_layer = GaussianAffinity(1, cfg.NGM.GAUSSIAN_SIGMA) 32 | else: 33 | raise ValueError('Unknown edge feature type {}'.format(cfg.NGM.EDGE_FEATURE)) 34 | self.tau = cfg.NGM.SK_TAU 35 | self.rescale = cfg.PROBLEM.RESCALE 36 | self.sinkhorn = Sinkhorn(max_iter=cfg.NGM.SK_ITER_NUM, tau=self.tau, epsilon=cfg.NGM.SK_EPSILON) 37 | self.gumbel_sinkhorn = GumbelSinkhorn(max_iter=cfg.NGM.SK_ITER_NUM, tau=self.tau * 10, epsilon=cfg.NGM.SK_EPSILON, batched_operation=True) 38 | self.l2norm = nn.LocalResponseNorm(cfg.NGM.FEATURE_CHANNEL * 2, alpha=cfg.NGM.FEATURE_CHANNEL * 2, beta=0.5, k=0) 39 | 40 | self.gnn_layer = cfg.NGM.GNN_LAYER 41 | for i in range(self.gnn_layer): 42 | tau = cfg.NGM.SK_TAU 43 | if i == 0: 44 | #gnn_layer = Gconv(1, cfg.NGM.GNN_FEAT) 45 | gnn_layer = GNNLayer(1, 1, cfg.NGM.GNN_FEAT[i] + (1 if cfg.NGM.SK_EMB else 0), cfg.NGM.GNN_FEAT[i], 46 | sk_channel=cfg.NGM.SK_EMB, sk_tau=tau, edge_emb=cfg.NGM.EDGE_EMB) 47 | #gnn_layer = HyperConvLayer(1, 1, cfg.NGM.GNN_FEAT[i] + (1 if cfg.NGM.SK_EMB else 0), cfg.NGM.GNN_FEAT[i], 48 | # sk_channel=cfg.NGM.SK_EMB, voting_alpha=alpha) 49 | else: 50 | #gnn_layer = Gconv(cfg.NGM.GNN_FEAT, cfg.NGM.GNN_FEAT) 51 | gnn_layer = GNNLayer(cfg.NGM.GNN_FEAT[i - 1] + (1 if cfg.NGM.SK_EMB else 0), cfg.NGM.GNN_FEAT[i - 1], 52 | cfg.NGM.GNN_FEAT[i] + (1 if cfg.NGM.SK_EMB else 0), cfg.NGM.GNN_FEAT[i], 53 | sk_channel=cfg.NGM.SK_EMB, sk_tau=tau, edge_emb=cfg.NGM.EDGE_EMB) 54 | #gnn_layer = HyperConvLayer(cfg.NGM.GNN_FEAT[i-1] + (1 if cfg.NGM.SK_EMB else 0), cfg.NGM.GNN_FEAT[i-1], 55 | # cfg.NGM.GNN_FEAT[i] + (1 if cfg.NGM.SK_EMB else 0), cfg.NGM.GNN_FEAT[i], 56 | # sk_channel=cfg.NGM.SK_EMB, voting_alpha=alpha) 57 | self.add_module('gnn_layer_{}'.format(i), gnn_layer) 58 | 59 | self.classifier = nn.Linear(cfg.NGM.GNN_FEAT[-1] + (1 if cfg.NGM.SK_EMB else 0), 1) 60 | 61 | def forward(self, data_dict, **kwargs): 62 | batch_size = data_dict['batch_size'] 63 | if 'images' in data_dict: 64 | # real image data 65 | src, tgt = data_dict['images'] 66 | P_src, P_tgt = data_dict['Ps'] 67 | ns_src, ns_tgt = data_dict['ns'] 68 | G_src, G_tgt = data_dict['Gs'] 69 | H_src, H_tgt = data_dict['Hs'] 70 | K_G, K_H = data_dict['KGHs'] 71 | 72 | # extract feature 73 | src_node = self.node_layers(src) 74 | src_edge = self.edge_layers(src_node) 75 | tgt_node = self.node_layers(tgt) 76 | tgt_edge = self.edge_layers(tgt_node) 77 | 78 | # feature normalization 79 | src_node = self.l2norm(src_node) 80 | src_edge = self.l2norm(src_edge) 81 | tgt_node = self.l2norm(tgt_node) 82 | tgt_edge = self.l2norm(tgt_edge) 83 | 84 | # arrange features 85 | U_src = feature_align(src_node, P_src, ns_src, self.rescale) 86 | F_src = feature_align(src_edge, P_src, ns_src, self.rescale) 87 | U_tgt = feature_align(tgt_node, P_tgt, ns_tgt, self.rescale) 88 | F_tgt = feature_align(tgt_edge, P_tgt, ns_tgt, self.rescale) 89 | elif 'features' in data_dict: 90 | # synthetic data 91 | src, tgt = data_dict['features'] 92 | P_src, P_tgt = data_dict['Ps'] 93 | ns_src, ns_tgt = data_dict['ns'] 94 | G_src, G_tgt = data_dict['Gs'] 95 | H_src, H_tgt = data_dict['Hs'] 96 | K_G, K_H = data_dict['KGHs'] 97 | 98 | U_src = src[:, :src.shape[1] // 2, :] 99 | F_src = src[:, src.shape[1] // 2:, :] 100 | U_tgt = tgt[:, :tgt.shape[1] // 2, :] 101 | F_tgt = tgt[:, tgt.shape[1] // 2:, :] 102 | elif 'aff_mat' in data_dict: 103 | K = data_dict['aff_mat'] 104 | ns_src, ns_tgt = data_dict['ns'] 105 | else: 106 | raise ValueError('Unknown data type for this model.') 107 | 108 | if 'images' in data_dict or 'features' in data_dict: 109 | tgt_len = P_tgt.shape[1] 110 | if cfg.NGM.EDGE_FEATURE == 'cat': 111 | X = reshape_edge_feature(F_src, G_src, H_src) 112 | Y = reshape_edge_feature(F_tgt, G_tgt, H_tgt) 113 | elif cfg.NGM.EDGE_FEATURE == 'geo': 114 | X = geo_edge_feature(P_src, G_src, H_src)[:, :1, :] 115 | Y = geo_edge_feature(P_tgt, G_tgt, H_tgt)[:, :1, :] 116 | else: 117 | raise ValueError('Unknown edge feature type {}'.format(cfg.NGM.EDGE_FEATURE)) 118 | 119 | # affinity layer 120 | Ke, Kp = self.affinity_layer(X, Y, U_src, U_tgt) 121 | 122 | K = construct_aff_mat(Ke, torch.zeros_like(Kp), K_G, K_H) 123 | 124 | A = (K > 0).to(K.dtype) 125 | 126 | if cfg.NGM.FIRST_ORDER: 127 | emb = Kp.transpose(1, 2).contiguous().view(Kp.shape[0], -1, 1) 128 | else: 129 | emb = torch.ones(K.shape[0], K.shape[1], 1, device=K.device) 130 | else: 131 | tgt_len = int(math.sqrt(K.shape[2])) 132 | dmax = (torch.max(torch.sum(K, dim=2, keepdim=True), dim=1, keepdim=True).values + 1e-5) 133 | K = K / dmax * 1000 134 | A = (K > 0).to(K.dtype) 135 | emb = torch.ones(K.shape[0], K.shape[1], 1, device=K.device) 136 | 137 | emb_K = K.unsqueeze(-1) 138 | 139 | # NGM qap solver 140 | for i in range(self.gnn_layer): 141 | gnn_layer = getattr(self, 'gnn_layer_{}'.format(i)) 142 | emb_K, emb = gnn_layer(A, emb_K, emb, ns_src, ns_tgt) #, norm=False) 143 | 144 | v = self.classifier(emb) 145 | s = v.view(v.shape[0], tgt_len, -1).transpose(1, 2) 146 | 147 | if self.training or cfg.NGM.GUMBEL_SK <= 0: 148 | #if cfg.NGM.GUMBEL_SK <= 0: 149 | ss = self.sinkhorn(s, ns_src, ns_tgt, dummy_row=True) 150 | x = hungarian(ss, ns_src, ns_tgt) 151 | else: 152 | gumbel_sample_num = cfg.NGM.GUMBEL_SK 153 | if self.training: 154 | gumbel_sample_num //= 10 155 | ss_gumbel = self.gumbel_sinkhorn(s, ns_src, ns_tgt, sample_num=gumbel_sample_num, dummy_row=True) 156 | 157 | repeat = lambda x, rep_num=gumbel_sample_num: torch.repeat_interleave(x, rep_num, dim=0) 158 | if not self.training: 159 | ss_gumbel = hungarian(ss_gumbel, repeat(ns_src), repeat(ns_tgt)) 160 | ss_gumbel = ss_gumbel.reshape(batch_size, gumbel_sample_num, ss_gumbel.shape[-2], ss_gumbel.shape[-1]) 161 | 162 | if ss_gumbel.device.type == 'cuda': 163 | dev_idx = ss_gumbel.device.index 164 | free_mem = gpu_free_memory(dev_idx) - 100 * 1024 ** 2 # 100MB as buffer for other computations 165 | K_mem_size = K.element_size() * K.nelement() 166 | max_repeats = free_mem // K_mem_size 167 | if max_repeats <= 0: 168 | print('Warning: GPU may not have enough memory') 169 | max_repeats = 1 170 | else: 171 | max_repeats = gumbel_sample_num 172 | 173 | obj_score = [] 174 | for idx in range(0, gumbel_sample_num, max_repeats): 175 | if idx + max_repeats > gumbel_sample_num: 176 | rep_num = gumbel_sample_num - idx 177 | else: 178 | rep_num = max_repeats 179 | obj_score.append( 180 | objective_score( 181 | ss_gumbel[:, idx:(idx+rep_num), :, :].reshape(-1, ss_gumbel.shape[-2], ss_gumbel.shape[-1]), 182 | repeat(K, rep_num) 183 | ).reshape(batch_size, -1) 184 | ) 185 | obj_score = torch.cat(obj_score, dim=1) 186 | min_obj_score = obj_score.min(dim=1) 187 | ss = ss_gumbel[torch.arange(batch_size), min_obj_score.indices.cpu(), :, :] 188 | x = hungarian(ss, repeat(ns_src), repeat(ns_tgt)) 189 | 190 | data_dict.update({ 191 | 'ds_mat': ss, 192 | 'perm_mat': x, 193 | 'aff_mat': K 194 | }) 195 | return data_dict 196 | --------------------------------------------------------------------------------