update: allow csq train at specified loss scale

Author: jiuntian

configs/templates/csq.yaml

@@ -2,4 +2,5 @@ arch: dpn
 loss: csq
 
 loss_param:
-  lambda_q: 0.001
+  lambda_q: 0.001
+  scale_c: 1.0

functions/loss/csq.py

@@ -6,11 +6,12 @@ class CSQLoss(nn.Module):
     """https://github.com/swuxyj/DeepHash-pytorch/blob/master/CSQ.py
     https://openaccess.thecvf.com/content_CVPR_2020/papers/Yuan_Central_Similarity_Quantization_for_Efficient_Image_and_Video_Retrieval_CVPR_2020_paper.pdf
     """
-    def __init__(self, multiclass, nbit, device, lambda_q=0.001, **kwargs):
+    def __init__(self, multiclass, nbit, device, lambda_q=0.001, scale_c=1., **kwargs):
         super(CSQLoss, self).__init__()
         device = torch.device(device)
         self.multiclass = multiclass
         self.lambda_q = lambda_q
+        self.scale_c = scale_c
         self.criterion = nn.BCELoss()
         self.multi_label_random_center = torch.randint(2, (nbit,)).float().to(device)
         self.losses = {}
@@ -27,7 +28,7 @@ def forward(self, logits, code_logits, labels, onehot=True):
         self.losses['center'] = loss_c
         self.losses['quant'] = loss_q
 
-        loss = loss_c + self.lambda_q * loss_q
+        loss = self.scale_c * loss_c + self.lambda_q * loss_q
         return loss
 
     def label2center(self, y, onehot):

scripts/train_general.py

@@ -123,7 +123,7 @@ def pre_epoch_operations(loss, **kwargs):
 
 
 def train_hashing(optimizer, model, train_loader, device, loss_name, loss_cfg, onehot,
-                  gpu_train_transform=None, method='supervised', criterion=None):
+                  gpu_train_transform=None, method='supervised', criterion=None, logdir=None):
     model.train()
 
     batch_timer = Timer()
@@ -133,7 +133,7 @@ def train_hashing(optimizer, model, train_loader, device, loss_name, loss_cfg, o
         criterion = train_helper.get_loss(loss_name, **loss_cfg)
     meters = defaultdict(AverageMeter)
 
-    train_helper.update_criterion(model=model, criterion=criterion, loss_name=loss_name)
+    train_helper.update_criterion(model=model, criterion=criterion, loss_name=loss_name, method=method, onehot=onehot)
     criterion.train()
 
     pbar = tqdm(train_loader, desc='Train', ascii=True, bar_format='{l_bar}{bar:10}{r_bar}',
@@ -182,6 +182,13 @@ def train_hashing(optimizer, model, train_loader, device, loss_name, loss_cfg, o
         running_times.append(batch_timer.total)
         pbar.set_postfix({key: val.avg for key, val in meters.items()})
         batch_timer.tick()
+
+        # if i % 2 == 0:
+        #     io.fast_save(output['code_logits'].detach().cpu(), f'{logdir}/outputs/train_iter_{i}.pth')
+        # if i > 200:
+        #     import sys
+        #     sys.exit(0)
+
     total_timer.toc()
     meters['total_time'].update(total_timer.total)
     std_time = f"time_std={np.std(running_times[1:]):.5f}"
@@ -206,7 +213,7 @@ def test_hashing(model, test_loader, device, loss_name, loss_cfg, onehot, return
     if criterion is None:
         criterion = train_helper.get_loss(loss_name, **loss_cfg)
 
-    train_helper.update_criterion(model=model, criterion=criterion, loss_name=loss_name)
+    train_helper.update_criterion(model=model, criterion=criterion, loss_name=loss_name, method=method, onehot=onehot)
     criterion.eval()
 
     pbar = tqdm(test_loader, desc='Test', ascii=True, bar_format='{l_bar}{bar:10}{r_bar}',
@@ -330,7 +337,9 @@ def preprocess(model, config, device):
         logging.info('Preprocessing for CSQ')
         nclass = config['arch_kwargs']['nclass']
         nbit = config['arch_kwargs']['nbit']
-        centroids = get_hadamard(nclass, nbit, fast=True)
+        # centroids = get_hadamard(nclass, nbit, fast=True)
+        centroids = generate_centroids(nclass, nbit, 'B')
+        logging.info("using bernoulli")
         centroids = centroids.to(device)
 
         # move to model
@@ -457,12 +466,6 @@ def main(config, gpu_transform=False, gpu_mean_transform=False, method='supervis
         ground_truth_path = os.path.join(test_loader.dataset.root, 'ground_truth.csv')
         ground_truth = pd.read_csv(ground_truth_path)  # id = index id, images = images id in database
 
-    # update criterion as non-onehot mode, for pairwise methods
-    if method in ['pairwise']:
-        if not onehot:
-            logging.info("Not a onehot label dataset")
-            criterion.label_not_onehot = True
-
     ##### resume training #####
     if config['start_epoch_from'] != 0:
         criterion, train_history, test_history = resume_training(config, logdir,
@@ -502,7 +505,7 @@ def main(config, gpu_transform=False, gpu_mean_transform=False, method='supervis
         train_meters = train_hashing(optimizer, model, train_loader, device, loss_param['loss'],
                                      loss_param['loss_param'], onehot=onehot,
                                      gpu_train_transform=gpu_train_transform,
-                                     method=method, criterion=criterion)
+                                     method=method, criterion=criterion, logdir=logdir)
 
         ##### scheduler #####
         if isinstance(scheduler, list):

scripts/train_helper.py

@@ -68,7 +68,7 @@ def get_loss(loss_name, **cfg):
     return loss[loss_name](**cfg)
 
 
-def update_criterion(model, criterion, loss_name):
+def update_criterion(model, criterion, loss_name, method, onehot):
     if loss_name in ['dpn', 'csq']:
         criterion.centroids = model.centroids
     elif loss_name in ['sdhc', 'sdh']:
@@ -79,6 +79,13 @@ def update_criterion(model, criterion, loss_name):
     elif loss_name in ['adsh']:
         criterion.weight = model.ce_fc.centroids
 
+    # update criterion as non-onehot mode, for pairwise methods
+    if method in ['pairwise']:
+
+        if not onehot and not criterion.label_not_onehot:
+            logging.info("Not a onehot label dataset")
+            criterion.label_not_onehot = True
+
 
 def generate_centroids(nclass, nbit, init_method):
     if init_method == 'N':