From e700c32731d67f47fa468224e9a82561f82986a8 Mon Sep 17 00:00:00 2001 From: KakaruHayate Date: Thu, 30 Jul 2026 17:22:23 +0800 Subject: [PATCH 1/2] refactor(training): clean up finetune/freeze code readability - reorganize into load_finetune_checkpoint + load_finetune_weights - use clear variable names (checkpoint_path, ignored_prefixes, mismatched_shapes) - replace per-key print loop with summarized rank_zero_info - avoid mutating input state dict in-place - simplify freeze_params with direct named_parameters iteration - delete dead code and dead comments from 2023-era - add actionable ValueError for missing/invalid checkpoint path - add weights_only=True for secure checkpoint loading --- basics/base_task.py | 132 ++++++++++++++++++++------------------------ 1 file changed, 60 insertions(+), 72 deletions(-) diff --git a/basics/base_task.py b/basics/base_task.py index 656893d96..02bca9bd7 100644 --- a/basics/base_task.py +++ b/basics/base_task.py @@ -89,92 +89,80 @@ def setup(self, stage): self.valid_dataset = self.dataset_cls('valid') self.num_replicas = (self.trainer.distributed_sampler_kwargs or {}).get('num_replicas', 1) - def get_need_freeze_state_dict_key(self, model_state_dict) -> list: - key_list = [] - for i in hparams['frozen_params']: - for j in model_state_dict: - if j.startswith(i): - key_list.append(j) - return list(set(key_list)) - def freeze_params(self) -> None: - model_state_dict = self.state_dict().keys() - freeze_key = self.get_need_freeze_state_dict_key(model_state_dict=model_state_dict) + frozen_prefixes = hparams['frozen_params'] + frozen_count = 0 + for name, parameter in self.named_parameters(): + if any(name.startswith(prefix) for prefix in frozen_prefixes): + parameter.requires_grad = False + frozen_count += 1 + rank_zero_info(f'| Frozen {frozen_count} parameter(s).') - for i in freeze_key: - params=self.get_parameter(i) + def unfreeze_all_params(self) -> None: + for parameter in self.model.parameters(): + parameter.requires_grad = True + + def load_finetune_weights(self, state_dict: dict) -> None: + source_state_dict = state_dict.copy() + if not hparams['finetune_strict_shapes']: + target_state_dict = self.state_dict() + mismatched_shapes = { + name: (tuple(source.shape), tuple(target_state_dict[name].shape)) + for name, source in source_state_dict.items() + if name in target_state_dict and source.shape != target_state_dict[name].shape + } + for name in mismatched_shapes: + del source_state_dict[name] + if mismatched_shapes: + details = '\n'.join( + f' {name}: source={source_shape}, target={target_shape}' + for name, (source_shape, target_shape) in mismatched_shapes.items() + ) + rank_zero_info( + f'| Finetune: skipped {len(mismatched_shapes)} parameter(s) with mismatched shapes:\n{details}' + ) + + incompatible_keys = self.load_state_dict(source_state_dict, strict=False) + loaded_count = len(source_state_dict) - len(incompatible_keys.unexpected_keys) + rank_zero_info( + f'| Finetune: loaded {loaded_count} parameter(s); ' + f'missing={len(incompatible_keys.missing_keys)}, ' + f'unexpected={len(incompatible_keys.unexpected_keys)}.' + ) - params.requires_grad = False + def load_finetune_checkpoint(self) -> dict: + checkpoint_path = hparams['finetune_ckpt_path'] + if not checkpoint_path: + raise ValueError('finetune_enabled=true requires a non-empty finetune_ckpt_path.') - def unfreeze_all_params(self) -> None: - for i in self.model.parameters(): - i.requires_grad = True - - def load_finetune_ckpt( - self, state_dict - ) -> None: - adapt_shapes = hparams['finetune_strict_shapes'] - if not adapt_shapes: - cur_model_state_dict = self.state_dict() - unmatched_keys = [] - for key, param in state_dict.items(): - if key in cur_model_state_dict: - new_param = cur_model_state_dict[key] - if new_param.shape != param.shape: - unmatched_keys.append(key) - print('| Unmatched keys: ', key, new_param.shape, param.shape) - for key in unmatched_keys: - del state_dict[key] - self.load_state_dict(state_dict, strict=False) - - def load_pre_train_model(self): - pre_train_ckpt_path = hparams['finetune_ckpt_path'] - blacklist = hparams['finetune_ignored_params'] - # whitelist=hparams['pre_train_whitelist'] - if blacklist is None: - blacklist = [] - # if whitelist is None: - # raise RuntimeError("") - - if pre_train_ckpt_path is not None: - ckpt = torch.load(pre_train_ckpt_path) - # if ckpt.get('category') is None: - # raise RuntimeError("") - - if isinstance(self.model, CategorizedModule): - self.model.check_category(ckpt.get('category')) - - state_dict = {} - for i in ckpt['state_dict']: - # if 'diffusion' in i: - # if i in rrrr: - # continue - skip = False - for b in blacklist: - if i.startswith(b): - skip = True - break - - if skip: - continue - - state_dict[i] = ckpt['state_dict'][i] - print(i) - return state_dict - else: - raise RuntimeError("") + checkpoint = torch.load(checkpoint_path, weights_only=True) + if isinstance(self.model, CategorizedModule): + self.model.check_category(checkpoint.get('category')) + + ignored_prefixes = hparams['finetune_ignored_params'] or [] + source_state_dict = checkpoint['state_dict'] + selected_state_dict = { + name: parameter + for name, parameter in source_state_dict.items() + if not any(name.startswith(prefix) for prefix in ignored_prefixes) + } + ignored_count = len(source_state_dict) - len(selected_state_dict) + rank_zero_info( + f"| Finetune: selected {len(selected_state_dict)} parameter(s) from '{checkpoint_path}'; " + f'ignored={ignored_count}.' + ) + return selected_state_dict def _build_model(self): raise NotImplementedError() def build_model(self): self.model = self._build_model() - # utils.load_warp(self) self.unfreeze_all_params() if hparams['freezing_enabled']: self.freeze_params() if hparams['finetune_enabled'] and get_latest_checkpoint_path(pathlib.Path(hparams['work_dir'])) is None: - self.load_finetune_ckpt(self.load_pre_train_model()) + self.load_finetune_weights(self.load_finetune_checkpoint()) self.print_arch() @rank_zero_only From ab567e9d161de353c3e4c4d930f60f5256288de5 Mon Sep 17 00:00:00 2001 From: KakaruHayate Date: Thu, 30 Jul 2026 18:15:23 +0800 Subject: [PATCH 2/2] fix(training): load finetune checkpoints on CPU --- basics/base_task.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/basics/base_task.py b/basics/base_task.py index 02bca9bd7..d83369e43 100644 --- a/basics/base_task.py +++ b/basics/base_task.py @@ -135,7 +135,7 @@ def load_finetune_checkpoint(self) -> dict: if not checkpoint_path: raise ValueError('finetune_enabled=true requires a non-empty finetune_ckpt_path.') - checkpoint = torch.load(checkpoint_path, weights_only=True) + checkpoint = torch.load(checkpoint_path, map_location='cpu', weights_only=True) if isinstance(self.model, CategorizedModule): self.model.check_category(checkpoint.get('category'))