File size: 3,852 Bytes
f15d29e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 | """Fine-tuning helpers backed by the example's local MatterGen model package."""
import logging
from collections import OrderedDict
from copy import deepcopy
from pathlib import Path
from typing import Tuple
import hydra
import pytorch_lightning as pl
import torch
from omegaconf import DictConfig, open_dict
from .common.utils.data_classes import MatterGenCheckpointInfo, _rewrite_vendored_targets
from .common.utils.globals import get_device
logger = logging.getLogger(__name__)
def init_adapter_lightningmodule_from_pretrained(
adapter_cfg: DictConfig, lightning_module_cfg: DictConfig
) -> Tuple[pl.LightningModule, DictConfig]:
"""Initialize the local adapter model from a MatterGen checkpoint."""
if adapter_cfg.model_path is not None:
if adapter_cfg.pretrained_name is not None:
logger.warning(
"pretrained_name is provided, but will be ignored since model_path is also provided."
)
model_path = Path(hydra.utils.to_absolute_path(adapter_cfg.model_path))
checkpoint_info = MatterGenCheckpointInfo(model_path, adapter_cfg.load_epoch)
elif adapter_cfg.pretrained_name is not None:
checkpoint_info = MatterGenCheckpointInfo.from_hf_hub(
adapter_cfg.pretrained_name
)
else:
raise ValueError("Either adapter.model_path or adapter.pretrained_name is required.")
checkpoint_path = checkpoint_info.checkpoint_path
version_root_path = Path(checkpoint_path).relative_to(
checkpoint_info.model_path
).parents[1]
config_path = Path(checkpoint_info.model_path) / version_root_path
pretrained_cfg_path = (
config_path if (config_path / "config.yaml").exists() else config_path.parent.parent
)
hydra.core.global_hydra.GlobalHydra.instance().clear()
with hydra.initialize_config_dir(
str(pretrained_cfg_path.absolute()), version_base="1.1"
):
pretrained_cfg = hydra.compose(config_name="config")
pretrained_cfg = _rewrite_vendored_targets(pretrained_cfg)
diffusion_module_cfg = deepcopy(pretrained_cfg.lightning_module.diffusion_module)
denoiser_cfg = diffusion_module_cfg.model
with open_dict(adapter_cfg.adapter):
for key, value in denoiser_cfg.items():
if key not in {"_target_", "property_embeddings_adapt"}:
adapter_cfg.adapter[key] = value
if key == "property_embeddings":
for field in value:
if field in adapter_cfg.adapter.property_embeddings_adapt:
adapter_cfg.adapter.property_embeddings_adapt.remove(field)
adapter_cfg.adapter.gemnet["_target_"] = (
"model.common.gemnet.gemnet_ctrl.GemNetTCtrl"
)
adapter_cfg.adapter.gemnet.condition_on_adapt = list(
adapter_cfg.adapter.property_embeddings_adapt
)
with open_dict(diffusion_module_cfg):
diffusion_module_cfg.model = adapter_cfg.adapter
with open_dict(lightning_module_cfg):
lightning_module_cfg.diffusion_module = diffusion_module_cfg
lightning_module = hydra.utils.instantiate(lightning_module_cfg)
checkpoint = torch.load(checkpoint_path, map_location=get_device())
pretrained_state: OrderedDict = checkpoint["state_dict"]
local_state: OrderedDict = lightning_module.state_dict()
local_state.update(
(key, pretrained_state[key])
for key in local_state.keys() & pretrained_state.keys()
)
lightning_module.load_state_dict(local_state, strict=True)
if not adapter_cfg.full_finetuning:
pretrained_keys = set(pretrained_state.keys())
for name, parameter in lightning_module.named_parameters():
if name in pretrained_keys:
parameter.requires_grad_(False)
return lightning_module, lightning_module_cfg
|