Rain-Bus/PromptSRC
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Last-k Average

Browse Source
This commit is contained in:
Rain&Bus
2026-02-07 15:58:51 +08:00
parent f3a7993665
commit 1d7d93ede5
7 changed files with 37 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files

63
trainers/promptsrc.py
View File

@@ -311,12 +311,8 @@ class PromptSRC(TrainerX):
# Cosine scheduler
self.total_epochs = cfg.OPTIM.MAX_EPOCH
self.step_counter = 1
N = cfg.OPTIM.MAX_EPOCH
mean = cfg.TRAINER.PROMPTSRC.GPA_MEAN
stdev = cfg.TRAINER.PROMPTSRC.GPA_STD
gauss = self.get_gauss(mean, stdev)
self.gauss = np.array([gauss(a) for a in range(1, N + 1)])
self.gauss = self.gauss / sum(self.gauss)
self.max_k = cfg.TRAINER.PROMPTSRC.LAST_K
self.last_k_models = []
self.scaler = GradScaler() if cfg.TRAINER.PROMPTSRC.PREC == "amp" else None
# Note that multi-gpu training could be slow because CLIP's size is
# big, which slows down the copy operation in DataParallel
@@ -324,8 +320,6 @@ class PromptSRC(TrainerX):
if device_count > 1:
print(f"Multiple GPUs detected (n_gpus={device_count}), use all of them!")
self.model = nn.DataParallel(self.model)
# Keep model with GPA
self.previous_model_gpa = None
def forward_backward(self, batch):
image, label = self.parse_batch_train(batch)
@@ -371,45 +365,32 @@ class PromptSRC(TrainerX):
if (self.batch_idx + 1) == self.num_batches:
self.update_lr()
# Means one epoch is completed, perform GPA
self.step_counter = self.step_counter + 1
current_epoch_weight = self.gauss[self.step_counter - 2]
current_model_weights = copy.deepcopy(model.state_dict())
weighted_state_dict = self.state_dict_weighting(current_model_weights, current_epoch_weight)
if self.previous_model_gpa is None:
self.previous_model_gpa = weighted_state_dict
else:
self.previous_model_gpa = self.state_dict_add(weighted_state_dict, self.previous_model_gpa)
for key in current_model_weights:
current_model_weights[key] = current_model_weights[key].cpu()
self.last_k_models.append(current_model_weights)
if len(self.last_k_models) > self.max_k:
self.last_k_models.pop(0)
torch.cuda.empty_cache()
if self.step_counter == self.model.total_epochs + 1:
print("Using GPA model for final inference...")
model.load_state_dict(self.previous_model_gpa)
self.model.load_state_dict(self.previous_model_gpa)
print(f"Using Last-K Averaging (K={len(self.last_k_models)}) model for final inference...")
averaged_state_dict = self._average_last_k_models()
for key in averaged_state_dict:
averaged_state_dict[key] = averaged_state_dict[key].cuda()
model.load_state_dict(averaged_state_dict)
self.model.load_state_dict(averaged_state_dict)
return loss_summary
def state_dict_weighting(self, main_dict, weightage, prompt_only=False):
# Average all parameters
updated_dict = copy.deepcopy(main_dict)
if not prompt_only:
for key in main_dict:
updated_dict[key] = main_dict[key] * weightage
return updated_dict
else:
return main_dict * weightage
def state_dict_add(self, dict1, dict2, prompt_only=False):
# Average all parameters
if not prompt_only:
modified_dict = dict2
for key in dict1:
modified_dict[key] = (modified_dict[key] + dict1[key])
return modified_dict
else:
return dict1 + dict2
def get_gauss(self, mu, sigma):
gauss = lambda x: (1 / (sigma * np.sqrt(2 * np.pi))) * np.exp(-0.5 * ((x - mu) / sigma) ** 2)
return gauss
def _average_last_k_models(self):
if not self.last_k_models:
return {}
averaged_dict = {}
for key in self.last_k_models[0]:
stacked = torch.stack([model_state[key] for model_state in self.last_k_models])
averaged_dict[key] = torch.mean(stacked, dim=0)
return averaged_dict
def parse_batch_train(self, batch):
input = batch["img"]