Method finetune_from

pixal3d/trainers/basic.py:489–554

Finetune from a checkpoint. Should be called by all processes.

(self, finetune_ckpt)

        return remapped_ckpt

    def finetune_from(self, finetune_ckpt):
        """
        Finetune from a checkpoint.
        Should be called by all processes.
        """
        # Allow missing keys (e.g., register_buffer parameters)
        ALLOWED_MISSING_KEYS = {'rope_phases'}
        
        if self.is_master:
            print('\nFinetuning from:')
            for name, path in finetune_ckpt.items():
                print(f'  - {name}: {path}')
        
        model_ckpts = {}
        for name, model in self.models.items():
            model_state_dict = model.state_dict()
            if name in finetune_ckpt:
                model_ckpt = torch.load(read_file_dist(finetune_ckpt[name]), map_location=self.device, weights_only=True)
                
                # Remap checkpoint keys to handle structural changes (e.g., ProjectAttention wrapper)
                model_ckpt = self._remap_checkpoint_keys(model_ckpt, model_state_dict)
                
                # Check extra keys (in ckpt but not in model)
                for k, v in model_ckpt.items():
                    if k not in model_state_dict:
                        if self.is_master:
                            print(f'Warning: {k} not found in model_state_dict, skipped.')
                        model_ckpt[k] = None
                    elif model_ckpt[k].shape != model_state_dict[k].shape:
                        if self.is_master:
                            print(f'Warning: {k} shape mismatch, {model_ckpt[k].shape} vs {model_state_dict[k].shape}, skipped.')
                        model_ckpt[k] = model_state_dict[k]
                model_ckpt = {k: v for k, v in model_ckpt.items() if v is not None}
                
                # Check missing keys (in model but not in ckpt)
                missing_keys = set(model_state_dict.keys()) - set(model_ckpt.keys())
                unexpected_missing = missing_keys - ALLOWED_MISSING_KEYS
                if unexpected_missing and self.is_master:
                    print(f'Error: Missing keys in checkpoint: {unexpected_missing}')
                    raise RuntimeError(f'Missing keys in checkpoint: {unexpected_missing}')
                if missing_keys & ALLOWED_MISSING_KEYS and self.is_master:
                    print(f'Info: Using model initialized values for: {missing_keys & ALLOWED_MISSING_KEYS}')
                
                # Fill in missing keys (using model initialized values)
                for k in missing_keys:
                    model_ckpt[k] = model_state_dict[k]
                
                model_ckpts[name] = model_ckpt
                model.load_state_dict(model_ckpt)
            else:
                if self.is_master:
                    print(f'Warning: {name} not found in finetune_ckpt, skipped.')
                model_ckpts[name] = model_state_dict
        self._state_dicts_to_master_params(self.master_params, model_ckpts)
        if self.is_master:
            for i, ema_rate in enumerate(self.ema_rate):
                self._state_dicts_to_master_params(self.ema_params[i], model_ckpts)
        del model_ckpts