File size: 4,702 Bytes
3e02ab8 | 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 95 96 97 98 99 100 101 102 103 104 105 106 107 108 | # This file is a part of the `nequip` package. Please see LICENSE and README at the root for information on using it.
import torch
from typing import Final, Callable, Optional, List
# NOTE: persistent modifiers are modifiers that fundamentally change the behavior of the model (same input will lead to different outputs)
# non-persistent modifiers generally refer to accelerations that should preserve similar model behavior, with the only difference being speed
_MODEL_MODIFIER_PERSISTENT_ATTR_NAME: Final[str] = (
"_nequip_model_modifier_is_persistent"
)
_MODEL_MODIFIER_PRIVATE_ATTR_NAME: Final[str] = "_nequip_model_modifier_is_private"
# these latter two attributes (unsupported devices and supported compile modes) are meant for acceleration modifiers
_MODEL_MODIFIER_UNSUPPORTED_DEVICES_ATTR_NAME: Final[str] = (
"_nequip_model_modifier_unsupported_devices"
)
_MODEL_MODIFIER_SUPPORTED_COMPILE_MODES_ATTR_NAME: Final[str] = (
"_nequip_model_modifier_supported_compile_modes"
)
def model_modifier(
persistent: bool,
private: Optional[bool] = None,
unsupported_devices: List[str] = [],
supported_compile_modes: Optional[List[str]] = None,
):
"""
Mark a ``@classmethod`` of an ``nn.Module`` as a "model modifier" that can be applied by the user to modify a packaged or other loaded model on-the-fly. Model modifiers must be a ``@classmethod`` of one of the ``nn.Module`` objects in the model.
Args:
persistent (bool): Whether the modifier should be applied when building the model for packaging.
private (bool, optional): Whether the modifier is private and should not be exposed in public interfaces. Defaults to None.
unsupported_devices (List[str], optional): List of device types that this modifier does not support. Defaults to [].
supported_compile_modes (List[str], optional): List of compile modes that this modifier supports. Defaults to None.
"""
def decorator(func):
assert isinstance(func, classmethod), (
"@model_modifier must be applied after @classmethod"
)
assert not hasattr(func.__func__, _MODEL_MODIFIER_PERSISTENT_ATTR_NAME)
setattr(func.__func__, _MODEL_MODIFIER_PERSISTENT_ATTR_NAME, persistent)
if private is not None:
setattr(func.__func__, _MODEL_MODIFIER_PRIVATE_ATTR_NAME, private)
setattr(
func.__func__,
_MODEL_MODIFIER_UNSUPPORTED_DEVICES_ATTR_NAME,
unsupported_devices,
)
setattr(
func.__func__,
_MODEL_MODIFIER_SUPPORTED_COMPILE_MODES_ATTR_NAME,
supported_compile_modes,
)
return func
return decorator
def is_model_modifier(func: callable) -> bool:
# for backwards compatibility, we use the "persistent" flag as a marker for whether the method is a model modifier
return hasattr(func, _MODEL_MODIFIER_PERSISTENT_ATTR_NAME)
def is_persistent_model_modifier(func: callable) -> bool:
return getattr(func, _MODEL_MODIFIER_PERSISTENT_ATTR_NAME)
def is_private_model_modifier(func: callable) -> Optional[bool]:
# for backwards compatibility of packaged models whose modifier would not have this metadata entry,
# we just default to making it public for convenience of clients
# should be ok since this mechanism is not safety critical and more just a convenience for documenting modifiers
return getattr(func, _MODEL_MODIFIER_PRIVATE_ATTR_NAME, False)
def get_model_modifier_unsupported_devices(func: callable) -> List[str]:
"""Get the list of unsupported devices for a model modifier. Returns empty list for backwards compatibility."""
return getattr(func, _MODEL_MODIFIER_UNSUPPORTED_DEVICES_ATTR_NAME, [])
def get_model_modifier_supported_compile_modes(func: callable) -> Optional[List[str]]:
"""Get the list of supported compile modes for a model modifier. Returns None if not set."""
return getattr(func, _MODEL_MODIFIER_SUPPORTED_COMPILE_MODES_ATTR_NAME, None)
def replace_submodules(
model: torch.nn.Module,
target_cls: type,
factory: Callable[[torch.nn.Module], torch.nn.Module],
) -> torch.nn.Module:
"""
Recursively walk the children of ``model``, and whenever we see an instance of ``target_cls``, replace it (in-place) with ``factory(old_module)`` by mutating ``model._modules[name]``.
"""
for name, child in list(model.named_children()):
if isinstance(child, target_cls):
# build a brand-new one based on `factory`
model._modules[name] = factory(child)
else:
# recurse down
replace_submodules(child, target_cls, factory)
return model
|