Personalised FedReID

.gitignore

@@ -0,0 +1,2 @@
+.idea/
+__pycache__/

client.py

@@ -8,7 +8,7 @@
 import copy
 from optimization import Optimization
 class Client():
-    def __init__(self, cid, data, device, project_dir, model_name, local_epoch, lr, batch_size, drop_rate, stride):
+    def __init__(self, cid, data, device, project_dir, model_name, local_epoch, lr, batch_size, drop_rate, stride, clustering=False):
         self.cid = cid
         self.project_dir = project_dir
         self.model_name = model_name
@@ -21,25 +21,30 @@ def __init__(self, cid, data, device, project_dir, model_name, local_epoch, lr,
         self.dataset_sizes = self.data.train_dataset_sizes[cid]
         self.train_loader = self.data.train_loaders[cid]
 
-        self.full_model = get_model(self.data.train_class_sizes[cid], drop_rate, stride)
-        self.classifier = self.full_model.classifier.classifier
-        self.full_model.classifier.classifier = nn.Sequential()
-        self.model = self.full_model
-        self.distance=0
+        self.model = get_model(self.data.train_class_sizes[cid], drop_rate, stride)
+        self.classifier = copy.deepcopy(self.model.classifier.classifier)
+        self.model.classifier.classifier = nn.Sequential()
+        self.distance = 0
         self.optimization = Optimization(self.train_loader, self.device)
+        self.use_clustering = clustering
         # print("class name size",class_names_size[cid])
 
-    def train(self, federated_model, use_cuda):
+    def train(self, federated_model=None, use_cuda=False):
         self.y_err = []
         self.y_loss = []
-
-        self.model.load_state_dict(federated_model.state_dict())
+        if self.use_clustering:
+            print("using clustering, model is set before")
+            assert federated_model is None
+            # self.model.classifier.classifier = nn.Sequential()
+            federated_model = copy.deepcopy(self.model)
+        else:
+            self.model.load_state_dict(federated_model.state_dict())
         self.model.classifier.classifier = self.classifier
         self.old_classifier = copy.deepcopy(self.classifier)
         self.model = self.model.to(self.device)
-
+        self.model.train(True)
         optimizer = get_optimizer(self.model, self.lr)
-        scheduler = lr_scheduler.StepLR(optimizer, step_size=40, gamma=0.1)
+        # scheduler = lr_scheduler.StepLR(optimizer, step_size=40, gamma=0.1)
 
         criterion = nn.CrossEntropyLoss()
 
@@ -50,8 +55,7 @@ def train(self, federated_model, use_cuda):
             print('Epoch {}/{}'.format(epoch, self.local_epoch - 1))
             print('-' * 10)
 
-            scheduler.step()
-            self.model.train(True)
+            # scheduler.step()
             running_loss = 0.0
             running_corrects = 0.0
 
@@ -60,12 +64,12 @@ def train(self, federated_model, use_cuda):
                 b, c, h, w = inputs.shape
                 if b < self.batch_size:
                     continue
-                if use_cuda:
-                    inputs = Variable(inputs.cuda().detach())
-                    labels = Variable(labels.cuda().detach())
-                else:
-                    inputs, labels = Variable(inputs), Variable(labels)
-
+                # if use_cuda:
+                #     inputs = Variable(inputs.cuda().detach())
+                #     labels = Variable(labels.cuda().detach())
+                # else:
+                #     inputs, labels = Variable(inputs), Variable(labels)
+                inputs, labels = inputs.to(self.device), labels.to(self.device)
                 optimizer.zero_grad()
 
                 outputs = self.model(inputs)
@@ -106,6 +110,9 @@ def train(self, federated_model, use_cuda):
     def generate_soft_label(self, x, regularization):
         return self.optimization.kd_generate_soft_label(self.model, x, regularization)
 
+    def generate_custom_data_feature(self, inputs):
+        return self.optimization.generate_custom_data_feature(self.model, inputs)
+
     def get_model(self):
         return self.model
 
@@ -116,4 +123,8 @@ def get_train_loss(self):
         return self.y_loss[-1]
 
     def get_cos_distance_weight(self):
-        return self.distance
+        return self.distance
+
+    def set_model(self, model):
+        self.model = copy.deepcopy(model)
+

data_utils.py

@@ -15,7 +15,7 @@ def __len__(self):
         return len(self.imgs)
 
     def __getitem__(self, index):
-        data,label = self.imgs[index]
+        data, label = self.imgs[index]
         return self.transform(Image.open(data)), label
 
 
@@ -95,6 +95,10 @@ def preprocess_train(self):
 
         print('Train dataset sizes:', self.train_dataset_sizes)
         print('Train class sizes:', self.train_class_sizes)
+        if "cuhk02" in self.datasets:
+            #cuhk02 is not labeled, we only use it for feature extraction in clustering
+            self.datasets.remove("cuhk02")
+            self.client_list.remove("cuhk02")
 
     def preprocess_test(self):
         """preprocess testing data, constructing test loaders
@@ -103,10 +107,8 @@ def preprocess_test(self):
         self.gallery_meta = {}
         self.query_meta = {}
 
-        for test_dir in self.datasets:
-            test_dir = 'data/'+test_dir+'/pytorch'
-
-            dataset = test_dir.split('/')[1]
+        for dataset in self.datasets:
+            test_dir = os.path.join(self.data_dir, dataset, "pytorch")
             gallery_dataset = datasets.ImageFolder(os.path.join(test_dir, 'gallery'))
             query_dataset = datasets.ImageFolder(os.path.join(test_dir, 'query'))
 

evaluate.py

@@ -2,11 +2,11 @@
 import torch
 import numpy as np
 import os
-import argparse
-parser = argparse.ArgumentParser(description='Training')
-parser.add_argument('--result_dir', default='.', type=str)
-parser.add_argument('--dataset', default='no_dataset', type=str)
-args = parser.parse_args()
+# import argparse
+# parser = argparse.ArgumentParser(description='Training')
+# parser.add_argument('--result_dir', default='.', type=str)
+# parser.add_argument('--dataset', default='no_dataset', type=str)
+# args = parser.parse_args()
 
 #######################################################################
 # Evaluate
@@ -60,32 +60,39 @@ def compute_mAP(index, good_index, junk_index):
 
     return ap, cmc
 
-######################################################################
-result = scipy.io.loadmat(args.result_dir + '/pytorch_result.mat')
 
-query_feature = torch.FloatTensor(result['query_f'])
-query_cam = result['query_cam'][0]
-query_label = result['query_label'][0]
-gallery_feature = torch.FloatTensor(result['gallery_f'])
-gallery_cam = result['gallery_cam'][0]
-gallery_label = result['gallery_label'][0]
+def testing_model(result, dataset):
+    # result = scipy.io.loadmat(file_path)
+    # print("========= after loading ==========")
+    # for i in result:
+    #     print(i, np.array(result[i]).shape)
 
-query_feature = query_feature.cuda()
-gallery_feature = gallery_feature.cuda()
+    query_feature = torch.FloatTensor(result['query_f'])
+    query_cam = np.array(result['query_cam'])
+    query_label = np.array(result['query_label'])
+    gallery_feature = torch.FloatTensor(result['gallery_f'])
+    gallery_cam = np.array(result['gallery_cam'])
+    gallery_label = np.array(result['gallery_label'])
+    # print(type(query_feature),query_feature[:3])
+    # print(type(query_cam),query_cam[:3])
+    # print(type(query_label),query_label[:3])
 
-print(query_feature.shape)
-CMC = torch.IntTensor(len(gallery_label)).zero_()
-ap = 0.0
+    query_feature = query_feature.cuda()
+    gallery_feature = gallery_feature.cuda()
 
-for i in range(len(query_label)):
-    ap_tmp, CMC_tmp = evaluate(query_feature[i], query_label[i], query_cam[i], gallery_feature, gallery_label, gallery_cam)
-    if CMC_tmp[0]==-1:
-        continue
-    CMC = CMC + CMC_tmp
-    ap += ap_tmp
+    print(query_feature.shape)
+    CMC = torch.IntTensor(len(gallery_label)).zero_()
+    ap = 0.0
 
-CMC = CMC.float()
-CMC = CMC/len(query_label) #average CMC
-print(args.dataset+' Rank@1:%f Rank@5:%f Rank@10:%f mAP:%f'%(CMC[0], CMC[4], CMC[9], ap/len(query_label)))
-print('-'*15)
-print()
+    for i in range(len(query_label)):
+        ap_tmp, CMC_tmp = evaluate(query_feature[i], query_label[i], query_cam[i], gallery_feature, gallery_label, gallery_cam)
+        if CMC_tmp[0]==-1:
+            continue
+        CMC = CMC + CMC_tmp
+        ap += ap_tmp
+
+    CMC = CMC.float()
+    CMC = CMC/len(query_label) #average CMC
+    print(dataset+' Rank@1:%f Rank@5:%f Rank@10:%f mAP:%f'%(CMC[0], CMC[4], CMC[9], ap/len(query_label)))
+    print('-'*15)
+    print()

finch.py

@@ -0,0 +1,173 @@
+import time
+import argparse
+import numpy as np
+from sklearn import metrics
+import scipy.sparse as sp
+import warnings
+
+try:
+    from pyflann import *
+
+    pyflann_available = True
+except Exception as e:
+    warnings.warn('pyflann not installed: {}'.format(e))
+    pyflann_available = False
+    pass
+
+RUN_FLANN = 70000
+
+
+def clust_rank(mat, initial_rank=None, distance='cosine'):
+    s = mat.shape[0]
+    if initial_rank is not None:
+        orig_dist = []
+    elif s <= RUN_FLANN:
+        orig_dist = metrics.pairwise.pairwise_distances(mat, mat, metric=distance)
+        np.fill_diagonal(orig_dist, 1000.0)
+        initial_rank = np.argmin(orig_dist, axis=1)
+    else:
+        if not pyflann_available:
+            raise MemoryError("You should use pyflann for inputs larger than {} samples.".format(RUN_FLANN))
+        print('Using flann to compute 1st-neighbours at this step ...')
+        flann = FLANN()
+        result, dists = flann.nn(mat, mat, num_neighbors=2, algorithm="kdtree", trees=8, checks=128)
+        initial_rank = result[:, 1]
+        orig_dist = []
+        print('Step flann done ...')
+
+    # The Clustering Equation
+    A = sp.csr_matrix((np.ones_like(initial_rank, dtype=np.float32), (np.arange(0, s), initial_rank)), shape=(s, s))
+    A = A + sp.eye(s, dtype=np.float32, format='csr')
+    A = A @ A.T
+
+    A = A.tolil()
+    A.setdiag(0)
+    return A, orig_dist
+
+
+def get_clust(a, orig_dist, min_sim=None):
+    if min_sim is not None:
+        a[np.where((orig_dist * a.toarray()) > min_sim)] = 0
+
+    num_clust, u = sp.csgraph.connected_components(csgraph=a, directed=True, connection='weak', return_labels=True)
+    return u, num_clust
+
+
+def cool_mean(M, u):
+    _, nf = np.unique(u, return_counts=True)
+    idx = np.argsort(u)
+    M = M[idx, :]
+    M = np.vstack((np.zeros((1, M.shape[1])), M))
+
+    np.cumsum(M, axis=0, out=M)
+    cnf = np.cumsum(nf)
+    nf1 = np.insert(cnf, 0, 0)
+    nf1 = nf1[:-1]
+
+    M = M[cnf, :] - M[nf1, :]
+    M = M / nf[:, None]
+    return M
+
+
+def get_merge(c, u, data):
+    if len(c) != 0:
+        _, ig = np.unique(c, return_inverse=True)
+        c = u[ig]
+    else:
+        c = u
+
+    mat = cool_mean(data, c)
+    return c, mat
+
+
+def update_adj(adj, d):
+    # Update adj, keep one merge at a time
+    idx = adj.nonzero()
+    v = np.argsort(d[idx])
+    v = v[:2]
+    x = [idx[0][v[0]], idx[0][v[1]]]
+    y = [idx[1][v[0]], idx[1][v[1]]]
+    a = sp.lil_matrix(adj.get_shape())
+    a[x, y] = 1
+    return a
+
+
+def req_numclust(c, data, req_clust, distance):
+    iter_ = len(np.unique(c)) - req_clust
+    c_, mat = get_merge([], c, data)
+    for i in range(iter_):
+        adj, orig_dist = clust_rank(mat, initial_rank=None, distance=distance)
+        adj = update_adj(adj, orig_dist)
+        u, _ = get_clust(adj, [], min_sim=None)
+        c_, mat = get_merge(c_, u, data)
+    return c_
+
+
+def FINCH(data, min_sim=None, initial_rank=None, req_clust=None, distance='cosine', verbose=True):
+    """ FINCH clustering algorithm.
+    :param data: Input matrix with features in rows.
+    :param initial_rank: Nx1 first integer neighbor indices (optional).
+    :param req_clust: Set output number of clusters (optional). Not recommended.
+    :param distance: One of ['cityblock', 'cosine', 'euclidean', 'l1', 'l2', 'manhattan'] Recommended 'cosine'.
+    :param verbose: Print verbose output.
+    :return:
+            c: NxP matrix where P is the partition. Cluster label for every partition.
+            num_clust: Number of clusters.
+            req_c: Labels of required clusters (Nx1). Only set if `req_clust` is not None.
+    The code implements the FINCH algorithm described in our CVPR 2019 paper
+        Sarfraz et al. "Efficient Parameter-free Clustering Using First Neighbor Relations", CVPR2019
+         https://arxiv.org/abs/1902.11266
+    For academic purpose only. The code or its re-implementation should not be used for commercial use.
+    Please contact the author below for licensing information.
+    Copyright
+    M. Saquib Sarfraz (saquib.sarfraz@kit.edu)
+    Karlsruhe Institute of Technology (KIT)
+    """
+    # Cast input data to float32
+    data = data.astype(np.float32)
+
+    # min_sim = None
+    adj, orig_dist = clust_rank(data, initial_rank, distance)
+    initial_rank = None
+    group, num_clust = get_clust(adj, [], min_sim)
+    c, mat = get_merge([], group, data)
+
+    if verbose:
+        print('Partition 0: {} clusters'.format(num_clust))
+    if len(orig_dist) != 0:
+        min_sim = np.max(orig_dist * adj.toarray())
+
+    exit_clust = 2
+    c_ = c
+    k = 1
+    num_clust = [num_clust]
+
+    while exit_clust > 1:
+        adj, orig_dist = clust_rank(mat, initial_rank, distance)
+        u, num_clust_curr = get_clust(adj, orig_dist, min_sim)
+        c_, mat = get_merge(c_, u, data)
+
+        num_clust.append(num_clust_curr)
+        c = np.column_stack((c, c_))
+        exit_clust = num_clust[-2] - num_clust_curr
+
+        if num_clust_curr == 1 or exit_clust < 1:
+            num_clust = num_clust[:-1]
+            c = c[:, :-1]
+            break
+
+        if verbose:
+            print('Partition {}: {} clusters'.format(k, num_clust[k]))
+        k += 1
+
+    if req_clust is not None:
+        if req_clust not in num_clust:
+            ind = [i for i, v in enumerate(num_clust) if v >= req_clust]
+            req_c = req_numclust(c[:, ind[-1]], data, req_clust, distance)
+        else:
+            req_c = c[:, num_clust.index(req_clust)]
+    else:
+        req_c = None
+
+    return c, num_clust, req_c
+