fix usam; add swa

Author: zzd@um4

model.py

@@ -27,7 +27,7 @@ def weights_init_classifier(m):
 
 class USAM(nn.Module):
     #Joint Representation Learning and Keypoint Detection for Cross-view Geo-localization. TIP2022
-    def __init__(self, kernel_size=3, padding=1, polish=True):
+    def __init__(self, kernel_size=3, padding=1, polish=False):
         super(USAM, self).__init__()
 
         kernel = torch.ones((kernel_size, kernel_size))
@@ -123,8 +123,8 @@ def __init__(self, class_num=751, droprate=0.5, stride=2, circle=False, ibn=Fals
         self.model = model_ft
         self.circle = circle
         self.classifier = ClassBlock(2048, class_num, droprate, linear=linear_num, return_f = circle)
+        self.usam = usam
         if usam:
-            self.usam = usam
             self.usam_1 = USAM()
             self.usam_2 = USAM()
 

test.py

@@ -17,6 +17,7 @@
 import scipy.io
 import yaml
 import math
+from torch.optim import swa_utils
 from tqdm import tqdm
 from model import ft_net, ft_net_dense, ft_net_hr, ft_net_swin, ft_net_swinv2, ft_net_efficient, ft_net_NAS, ft_net_convnext, PCB, PCB_test
 from utils import fuse_all_conv_bn
@@ -167,9 +168,15 @@ def load_network(network):
             print("Compiling model...")
             # https://huggingface.co/docs/diffusers/main/en/optimization/torch2.0
             torch.set_float32_matmul_precision('high')
-            network = torch.compile(network, mode="default", dynamic=True) # pytorch 2.0
+            network.cuda()
+            network = torch.compile(network, mode="reduce-overhead", dynamic = True) # pytorch 2.0
+        if 'average' in opt.which_epoch: # load averaged model.
+            network = swa_utils.AveragedModel(network)
         network.load_state_dict(torch.load(save_path))
-
+        if 'average' in opt.which_epoch:
+            print("We average %d snapshots"%network.n_averaged)
+            #swa_utils.update_bn(dataloaders['query'], network, device='cuda:0')
+            network = network.module
     return network
 
 

train.py

@@ -16,6 +16,7 @@
 import time
 import os
 import collections
+from torch.optim import swa_utils
 from tqdm import tqdm
 from model import ft_net, ft_net_dense, ft_net_hr, ft_net_swin, ft_net_swinv2, ft_net_convnext, ft_net_efficient, ft_net_NAS, PCB
 from random_erasing import RandomErasing
@@ -57,6 +58,7 @@
 parser.add_argument('--fp16', action='store_true', help='use float16 instead of float32, which will save about 50%% memory' )
 parser.add_argument('--cosine', action='store_true', help='use cosine lrRate' )
 parser.add_argument('--FSGD', action='store_true', help='use fused sgd, which will speed up trainig slightly. apex is needed.' )
+parser.add_argument('--wa', action='store_true', help='use weight average' )
 # backbone
 parser.add_argument('--linear_num', default=512, type=int, help='feature dimension: 512 or default or 0 (linear=False)')
 parser.add_argument('--stride', default=2, type=int, help='stride')
@@ -88,6 +90,9 @@
 
 opt = parser.parse_args()
 
+if opt.DG:
+    opt.wa = True #DG will enable swa.
+
 fp16 = opt.fp16
 data_dir = opt.data_dir
 name = opt.name
@@ -221,6 +226,7 @@ def train_model(model, criterion, optimizer, scheduler, num_epochs=25):
 
     #best_model_wts = model.state_dict()
     #best_acc = 0.0
+    wa_flag = opt.wa
     warm_up = 0.1 # We start from the 0.1*lrRate
     warm_iteration = round(dataset_sizes['train']/opt.batchsize)*opt.warm_epoch # first 5 epoch
     embedding_size = model.classifier.linear_num
@@ -244,6 +250,12 @@ def train_model(model, criterion, optimizer, scheduler, num_epochs=25):
     for epoch in range(num_epochs):
         print('Epoch {}/{}'.format(epoch, num_epochs - 1))
         # print('-' * 10)
+
+        if opt.wa and wa_flag and epoch >=  num_epochs*0.1:
+            wa_flag = False
+            swa_model = swa_utils.AveragedModel(model)
+            swa_model.avg_fn = swa_utils.get_ema_avg_fn(decay=0.996)
+            print('start weight avg')
 
         # Each epoch has a training and validation phase
         for phase in ['train', 'val']:
@@ -286,8 +298,6 @@ def train_model(model, criterion, optimizer, scheduler, num_epochs=25):
                 else:
                     outputs = model(inputs)
 
-
-
                 if opt.adv>0 and iter%opt.aiter==0: 
                     inputs_adv = ODFA(model, inputs)
                     outputs_adv = model(inputs_adv)
@@ -365,17 +375,22 @@ def train_model(model, criterion, optimizer, scheduler, num_epochs=25):
                         for i in range(num_part):
                             part[i] = outputs1[i]
                         outputs1 = part[0] + part[1] + part[2] + part[3] + part[4] + part[5]
-                    outputs2 = model(inputs2)
+
+                    swa_model.eval()
+                    with torch.no_grad():
+                        outputs2 = swa_model(inputs2) #stop gradient like dino
+                    outputs2 = outputs2.detach()
+
                     if return_feature:
                         outputs2, _ = outputs2
                     elif opt.PCB:
                         for i in range(num_part):
                             part[i] = outputs2[i]
                         outputs2 = part[0] + part[1] + part[2] + part[3] + part[4] + part[5]
 
-                    mean_pred = sm(outputs1 + outputs2)
+                    #supervised via teacher like dino. previous use sm(outputs1 + outputs2)
                     kl_loss = nn.KLDivLoss(reduction='batchmean')
-                    reg= (kl_loss(log_sm(outputs2) , mean_pred)  + kl_loss(log_sm(outputs1) , mean_pred))/2
+                    reg= (kl_loss(log_sm(outputs2), sm(outputs1))  + kl_loss(log_sm(outputs1) , sm(outputs2)))/2
                     loss += 0.01*reg
                     del inputs1, inputs2
                     #print(0.01*reg)
@@ -419,6 +434,10 @@ def train_model(model, criterion, optimizer, scheduler, num_epochs=25):
             pbar.set_postfix(ordered_dict=ordered_dict)
             pbar.close()
 
+            if phase == 'train' and opt.wa and epoch >= num_epochs*0.1: 
+                swa_model.update_parameters(model)
+                swa_utils.update_bn(dataloaders['train'], swa_model, device='cuda:0')
+
             y_loss[phase].append(epoch_loss)
             y_err[phase].append(1.0-epoch_acc)            
             # deep copy the model
@@ -449,6 +468,11 @@ def train_model(model, criterion, optimizer, scheduler, num_epochs=25):
     else:
         save_network(model, opt.name, 'last')
 
+    if opt.wa:
+         save_network( swa_model, opt.name, 'average')
+         swa_utils.update_bn(dataloaders['train'], swa_model, device='cuda:0')
+         save_network( swa_model, opt.name, 'average_bn')
+
     return model
 
 
@@ -511,6 +535,7 @@ def draw_curve(current_epoch):
 
 if torch.cuda.get_device_capability()[0]>6 and len(opt.gpu_ids)==1 and int(version[0])>1: # should be >=7 and one gpu
     torch.set_float32_matmul_precision('high')
+    torch._dynamo.config.automatic_dynamic_shapes = True
     print("Compiling model... The first epoch may be slow, which is expected!")
     # https://huggingface.co/docs/diffusers/main/en/optimization/torch2.0
     model = torch.compile(model, mode="reduce-overhead", dynamic = True) # pytorch 2.0