|
| 1 | +# SPDX-License-Identifier: LGPL-3.0-or-later |
| 2 | +"""Backend-neutral helpers for multi-task shared-parameter wiring.""" |
| 3 | + |
| 4 | +import logging |
| 5 | +from collections.abc import ( |
| 6 | + Callable, |
| 7 | +) |
| 8 | +from copy import ( |
| 9 | + deepcopy, |
| 10 | +) |
| 11 | +from typing import ( |
| 12 | + Any, |
| 13 | +) |
| 14 | + |
| 15 | +log = logging.getLogger(__name__) |
| 16 | + |
| 17 | + |
| 18 | +def cascade_top_level_defaults(model_config: dict[str, Any]) -> None: |
| 19 | + """Lower model-wide entries into each multi-task branch in-place.""" |
| 20 | + reserved_top_level = ("model_dict", "shared_dict") |
| 21 | + top_level_defaults = { |
| 22 | + k: deepcopy(v) for k, v in model_config.items() if k not in reserved_top_level |
| 23 | + } |
| 24 | + for branch in model_config["model_dict"].values(): |
| 25 | + for k, v in top_level_defaults.items(): |
| 26 | + branch.setdefault(k, deepcopy(v)) |
| 27 | + for k in top_level_defaults: |
| 28 | + model_config.pop(k, None) |
| 29 | + |
| 30 | + |
| 31 | +def preprocess_shared_params( |
| 32 | + model_config: dict[str, Any], |
| 33 | + get_class_name: Callable[[str, dict[str, Any]], type], |
| 34 | + *, |
| 35 | + require_shared_type_map: bool = True, |
| 36 | + cascade_defaults: bool = False, |
| 37 | +) -> tuple[dict[str, Any], dict[str, Any]]: |
| 38 | + """Expand ``shared_dict`` references and generate runtime sharing links. |
| 39 | +
|
| 40 | + Parameters |
| 41 | + ---------- |
| 42 | + model_config : dict[str, Any] |
| 43 | + Multi-task model configuration containing ``model_dict`` and an optional |
| 44 | + ``shared_dict``. String references to shared descriptors, fitting |
| 45 | + networks, and type maps are replaced in-place by their configurations. |
| 46 | + get_class_name : Callable[[str, dict[str, Any]], type] |
| 47 | + Backend-specific callback that resolves the class for a shared |
| 48 | + descriptor or fitting-network configuration. |
| 49 | + require_shared_type_map : bool, default=True |
| 50 | + Whether exactly one shared type-map reference is required. If false, at |
| 51 | + most one shared type-map reference is allowed. |
| 52 | + cascade_defaults : bool, default=False |
| 53 | + Whether to copy non-reserved top-level model options into each model |
| 54 | + branch that does not define them, then remove those options from the top |
| 55 | + level. |
| 56 | +
|
| 57 | + Returns |
| 58 | + ------- |
| 59 | + model_config : dict[str, Any] |
| 60 | + The input configuration with shared references expanded. |
| 61 | + shared_links : dict[str, Any] |
| 62 | + Sharing metadata keyed by entries in ``shared_dict``. Each entry records |
| 63 | + the resolved class and the model components linked to it, sorted by |
| 64 | + sharing level. |
| 65 | + """ |
| 66 | + assert "model_dict" in model_config, "only multi-task model can use this method!" |
| 67 | + if cascade_defaults: |
| 68 | + cascade_top_level_defaults(model_config) |
| 69 | + |
| 70 | + supported_types = ["type_map", "descriptor", "fitting_net"] |
| 71 | + shared_dict = model_config.get("shared_dict", {}) |
| 72 | + shared_links: dict[str, Any] = {} |
| 73 | + type_map_keys: list[str] = [] |
| 74 | + |
| 75 | + def replace_one_item( |
| 76 | + params_dict: dict[str, Any] | list[Any], |
| 77 | + model_key: str, |
| 78 | + key_type: str, |
| 79 | + key_in_dict: str, |
| 80 | + suffix: str = "", |
| 81 | + index: int | None = None, |
| 82 | + ) -> None: |
| 83 | + shared_type = key_type |
| 84 | + shared_key = key_in_dict |
| 85 | + shared_level = 0 |
| 86 | + if ":" in key_in_dict: |
| 87 | + shared_key = key_in_dict.split(":")[0] |
| 88 | + shared_level = int(key_in_dict.split(":")[1]) |
| 89 | + assert shared_key in shared_dict, ( |
| 90 | + f"Appointed {shared_type} {shared_key} are not in the shared_dict! " |
| 91 | + "Please check the input params." |
| 92 | + ) |
| 93 | + if index is None: |
| 94 | + assert isinstance(params_dict, dict) |
| 95 | + params_dict[shared_type] = deepcopy(shared_dict[shared_key]) |
| 96 | + else: |
| 97 | + params_dict[index] = deepcopy(shared_dict[shared_key]) |
| 98 | + if shared_type == "type_map": |
| 99 | + if key_in_dict not in type_map_keys: |
| 100 | + type_map_keys.append(key_in_dict) |
| 101 | + else: |
| 102 | + if shared_key not in shared_links: |
| 103 | + shared_links[shared_key] = { |
| 104 | + "type": get_class_name(shared_type, shared_dict[shared_key]), |
| 105 | + "links": [], |
| 106 | + } |
| 107 | + shared_links[shared_key]["links"].append( |
| 108 | + { |
| 109 | + "model_key": model_key, |
| 110 | + "shared_type": shared_type + suffix, |
| 111 | + "shared_level": shared_level, |
| 112 | + } |
| 113 | + ) |
| 114 | + |
| 115 | + for model_key in model_config["model_dict"]: |
| 116 | + model_params_item = model_config["model_dict"][model_key] |
| 117 | + for item_key in model_params_item: |
| 118 | + if item_key not in supported_types: |
| 119 | + continue |
| 120 | + item_params = model_params_item[item_key] |
| 121 | + if isinstance(item_params, str): |
| 122 | + replace_one_item(model_params_item, model_key, item_key, item_params) |
| 123 | + elif ( |
| 124 | + isinstance(item_params, dict) |
| 125 | + and item_params.get("type", "") == "hybrid" |
| 126 | + ): |
| 127 | + for ii, hybrid_item in enumerate(item_params["list"]): |
| 128 | + if isinstance(hybrid_item, str): |
| 129 | + replace_one_item( |
| 130 | + model_params_item[item_key]["list"], |
| 131 | + model_key, |
| 132 | + item_key, |
| 133 | + hybrid_item, |
| 134 | + suffix=f"_hybrid_{ii}", |
| 135 | + index=ii, |
| 136 | + ) |
| 137 | + |
| 138 | + for shared_key in shared_links: |
| 139 | + shared_links[shared_key]["links"] = sorted( |
| 140 | + shared_links[shared_key]["links"], |
| 141 | + key=lambda x: ( |
| 142 | + x["shared_level"] |
| 143 | + - ("spin" in model_config["model_dict"][x["model_key"]]) * 100 |
| 144 | + ), |
| 145 | + ) |
| 146 | + if require_shared_type_map: |
| 147 | + assert len(type_map_keys) == 1, "Multitask model must have only one type_map!" |
| 148 | + else: |
| 149 | + assert len(type_map_keys) <= 1, "Shared params must have at most one type_map!" |
| 150 | + return model_config, shared_links |
| 151 | + |
| 152 | + |
| 153 | +def apply_shared_links( |
| 154 | + models: dict[str, Any], |
| 155 | + shared_links: dict[str, Any] | None, |
| 156 | + *, |
| 157 | + share_descriptor: Callable[..., None], |
| 158 | + share_fitting: Callable[..., None], |
| 159 | + model_key_prob_map: dict[str, float] | None = None, |
| 160 | + data_stat_protect: float = 1e-2, |
| 161 | + resume: bool = False, |
| 162 | + get_descriptor: Callable[[Any, str], Any] | None = None, |
| 163 | + get_fitting: Callable[[Any, str], Any | None] | None = None, |
| 164 | + logger: logging.Logger | None = None, |
| 165 | +) -> None: |
| 166 | + """Apply sharing links to constructed models with backend-specific hooks. |
| 167 | +
|
| 168 | + Parameters |
| 169 | + ---------- |
| 170 | + models : dict[str, Any] |
| 171 | + Constructed models keyed by the model names used in ``shared_links``. |
| 172 | + shared_links : dict[str, Any] or None |
| 173 | + Sharing metadata generated by :func:`preprocess_shared_params`. No work |
| 174 | + is performed when it is empty or ``None``. |
| 175 | + share_descriptor : Callable[..., None] |
| 176 | + Backend-specific callback that shares a linked descriptor with its base |
| 177 | + descriptor. |
| 178 | + share_fitting : Callable[..., None] |
| 179 | + Backend-specific callback that shares a linked fitting network with its |
| 180 | + base fitting network. |
| 181 | + model_key_prob_map : dict[str, float] or None, default=None |
| 182 | + Sampling probabilities keyed by model name. The linked-to-base |
| 183 | + probability ratio is passed to the sharing callbacks to weight merged |
| 184 | + statistics. Equal probabilities are used when omitted. |
| 185 | + data_stat_protect : float, default=1e-2 |
| 186 | + Protection value passed to the fitting-network sharing callback when |
| 187 | + statistics are merged. |
| 188 | + resume : bool, default=False |
| 189 | + Whether sharing is being restored from a checkpoint. Backend callbacks |
| 190 | + use this flag to skip merging statistics again. |
| 191 | + get_descriptor : Callable[[Any, str], Any] or None, default=None |
| 192 | + Optional accessor for a descriptor or hybrid sub-descriptor. The common |
| 193 | + model accessor is used when omitted. |
| 194 | + get_fitting : Callable[[Any, str], Any or None] or None, default=None |
| 195 | + Optional accessor for a fitting component. The common atomic-model |
| 196 | + accessor is used when omitted. |
| 197 | + logger : logging.Logger or None, default=None |
| 198 | + Logger used to report parameter-sharing operations. |
| 199 | +
|
| 200 | + Returns |
| 201 | + ------- |
| 202 | + None |
| 203 | + The components in ``models`` are shared in-place. |
| 204 | + """ |
| 205 | + if not shared_links: |
| 206 | + return |
| 207 | + if model_key_prob_map is None: |
| 208 | + model_key_prob_map = dict.fromkeys(models, 1.0) |
| 209 | + if get_descriptor is None: |
| 210 | + get_descriptor = get_descriptor_component |
| 211 | + if get_fitting is None: |
| 212 | + get_fitting = get_atomic_model_attr |
| 213 | + if logger is None: |
| 214 | + logger = log |
| 215 | + |
| 216 | + for shared_info in shared_links.values(): |
| 217 | + links = shared_info.get("links", []) |
| 218 | + if len(links) < 2: |
| 219 | + continue |
| 220 | + shared_base = links[0] |
| 221 | + class_type_base = shared_base["shared_type"] |
| 222 | + model_key_base = shared_base["model_key"] |
| 223 | + shared_level_base = int(shared_base["shared_level"]) |
| 224 | + |
| 225 | + if "descriptor" in class_type_base: |
| 226 | + base_class = get_descriptor(models[model_key_base], class_type_base) |
| 227 | + for link_item in links[1:]: |
| 228 | + class_type_link = link_item["shared_type"] |
| 229 | + model_key_link = link_item["model_key"] |
| 230 | + shared_level_link = int(link_item["shared_level"]) |
| 231 | + assert shared_level_link >= shared_level_base, ( |
| 232 | + "The shared_links must be sorted by shared_level!" |
| 233 | + ) |
| 234 | + assert "descriptor" in class_type_link, ( |
| 235 | + f"Class type mismatched: {class_type_base} vs {class_type_link}!" |
| 236 | + ) |
| 237 | + link_model = models[model_key_link] |
| 238 | + link_class = get_descriptor(link_model, class_type_link) |
| 239 | + share_descriptor( |
| 240 | + link_model, |
| 241 | + class_type_link, |
| 242 | + link_class, |
| 243 | + base_class, |
| 244 | + shared_level_link, |
| 245 | + _model_prob_ratio( |
| 246 | + model_key_prob_map, model_key_base, model_key_link |
| 247 | + ), |
| 248 | + resume=resume, |
| 249 | + ) |
| 250 | + logger.warning( |
| 251 | + "Shared params of %s.%s and %s.%s!", |
| 252 | + model_key_base, |
| 253 | + class_type_base, |
| 254 | + model_key_link, |
| 255 | + class_type_link, |
| 256 | + ) |
| 257 | + continue |
| 258 | + |
| 259 | + base_class = get_fitting(models[model_key_base], class_type_base) |
| 260 | + if base_class is None: |
| 261 | + continue |
| 262 | + for link_item in links[1:]: |
| 263 | + class_type_link = link_item["shared_type"] |
| 264 | + model_key_link = link_item["model_key"] |
| 265 | + shared_level_link = int(link_item["shared_level"]) |
| 266 | + assert shared_level_link >= shared_level_base, ( |
| 267 | + "The shared_links must be sorted by shared_level!" |
| 268 | + ) |
| 269 | + assert class_type_base == class_type_link, ( |
| 270 | + f"Class type mismatched: {class_type_base} vs {class_type_link}!" |
| 271 | + ) |
| 272 | + link_class = get_fitting(models[model_key_link], class_type_link) |
| 273 | + if link_class is None: |
| 274 | + continue |
| 275 | + share_fitting( |
| 276 | + link_class, |
| 277 | + base_class, |
| 278 | + shared_level_link, |
| 279 | + _model_prob_ratio(model_key_prob_map, model_key_base, model_key_link), |
| 280 | + protection=data_stat_protect, |
| 281 | + resume=resume, |
| 282 | + ) |
| 283 | + logger.warning( |
| 284 | + "Shared params of %s.%s and %s.%s!", |
| 285 | + model_key_base, |
| 286 | + class_type_base, |
| 287 | + model_key_link, |
| 288 | + class_type_link, |
| 289 | + ) |
| 290 | + |
| 291 | + |
| 292 | +def get_descriptor_component(model: Any, shared_type: str) -> Any: |
| 293 | + """Get a model descriptor or a hybrid sub-descriptor by shared type.""" |
| 294 | + if shared_type == "descriptor": |
| 295 | + return model.get_descriptor() |
| 296 | + if "hybrid" in shared_type: |
| 297 | + hybrid_index = int(shared_type.split("_")[-1]) |
| 298 | + return model.get_descriptor().descrpt_list[hybrid_index] |
| 299 | + raise RuntimeError(f"Unknown class_type {shared_type}!") |
| 300 | + |
| 301 | + |
| 302 | +def set_descriptor_component(model: Any, shared_type: str, descriptor: Any) -> None: |
| 303 | + """Set a model descriptor or a hybrid sub-descriptor by shared type.""" |
| 304 | + if shared_type == "descriptor": |
| 305 | + model.atomic_model.descriptor = descriptor |
| 306 | + return |
| 307 | + if "hybrid" in shared_type: |
| 308 | + hybrid_index = int(shared_type.split("_")[-1]) |
| 309 | + model.get_descriptor().descrpt_list[hybrid_index] = descriptor |
| 310 | + return |
| 311 | + raise RuntimeError(f"Unknown class_type {shared_type}!") |
| 312 | + |
| 313 | + |
| 314 | +def get_atomic_model_attr(model: Any, attr: str) -> Any | None: |
| 315 | + """Get an attribute from ``model.atomic_model`` if it exists.""" |
| 316 | + if hasattr(model.atomic_model, attr): |
| 317 | + return getattr(model.atomic_model, attr) |
| 318 | + return None |
| 319 | + |
| 320 | + |
| 321 | +def sanitize_shared_links(shared_links: dict[str, Any] | None) -> dict[str, Any] | None: |
| 322 | + """Return a JSON-safe copy of ``shared_links``.""" |
| 323 | + if shared_links is None: |
| 324 | + return None |
| 325 | + sanitized: dict[str, Any] = {} |
| 326 | + for shared_key, shared_info in shared_links.items(): |
| 327 | + class_type = shared_info.get("type") |
| 328 | + sanitized[shared_key] = { |
| 329 | + "type": getattr(class_type, "__name__", str(class_type)), |
| 330 | + "links": deepcopy(shared_info.get("links", [])), |
| 331 | + } |
| 332 | + return sanitized |
| 333 | + |
| 334 | + |
| 335 | +def _model_prob_ratio( |
| 336 | + model_key_prob_map: dict[str, float], |
| 337 | + model_key_base: str, |
| 338 | + model_key_link: str, |
| 339 | +) -> float: |
| 340 | + return float(model_key_prob_map[model_key_link]) / float( |
| 341 | + model_key_prob_map[model_key_base] |
| 342 | + ) |
0 commit comments