fusion-bench 0.2.8__py3-none-any.whl → 0.2.10__py3-none-any.whl

fusion_bench/utils/expr.py

@@ -0,0 +1,90 @@
+class pattern_query:
+    R"""
+    Examples:
+
+        >>> f = pattern_query(lambda x: x==1)
+        >>> f(1)
+        True
+        >>> f(2)
+        False
+
+    """
+
+    def __init__(self, func=None, type=None):
+        self.func = func
+        self.type = type
+
+    def __call__(self, expr) -> bool:
+        if self.type is not None:
+            if not isinstance(expr, self.type):
+                return False
+        if self.func is not None:
+            return self.func(expr)
+        return True
+
+
+def is_expr_match(pattern, expr):
+    R"""match pattern with a python expression expr.
+
+    Examples:
+
+        >>> is_expr_match('a', 'a')
+        True
+        >>> is_expr_match((object, 1), ('s',1))
+        True
+        >>> is_expr_match((object, 1), ('s',2))
+        False
+        >>> is_expr_match(((int, (int,)), (int, (int,)), (-1,)),
+                          ((2146, (6,)), (1124, (97,)), (-1,)))
+        True
+
+        match a numpy array whose shape is (1,2)
+
+        >>> import numpy as np
+        >>> is_expr_match(
+                pattern_query(lambda arr: arr.shape==(1,2), np.ndarray),
+                np.zeros((1,2)))
+        True
+
+    Args:
+        pattern: pattern to match
+        expr:  python object
+
+    Raises:
+        NotImplementedError: Unsupported type
+
+    Returns:
+        bool
+    """
+
+    if type(pattern) == type:  # type
+        if not isinstance(expr, pattern):
+            return False
+
+    else:  # instance
+        if isinstance(pattern, pattern_query):
+            return pattern(expr)
+        if type(pattern) != type(expr):
+            return False
+
+        if isinstance(pattern, (int, float, str)):
+            if pattern != expr:
+                return False
+        elif isinstance(pattern, (tuple, list)):
+            if len(pattern) != len(expr):
+                return False
+            for i in range(len(pattern)):
+                if not is_expr_match(pattern[i], expr[i]):
+                    return False
+        elif isinstance(pattern, dict):
+            if len(pattern) != len(expr):
+                return False
+            for k in pattern:
+                try:
+                    if not is_expr_match(pattern[k], expr[k]):
+                        return False
+                except:
+                    return False
+        else:
+            raise NotImplementedError("Unsupported type: {}".format(type(pattern)))
+    return True

fusion_bench/utils/fabric.py

@@ -0,0 +1,17 @@
+import time
+
+import lightning as L
+
+
+def seed_everything_by_time(fabric: L.Fabric):
+    """
+    Set seed for all processes by time.
+    """
+    # set seed for all processes
+    if fabric.is_global_zero:
+        seed = int(time.time())
+    else:
+        seed = None
+    fabric.barrier()
+    seed = fabric.broadcast(seed, src=0)
+    L.seed_everything(seed)

fusion_bench/utils/instantiate.py

@@ -10,7 +10,7 @@ from hydra._internal.utils import _locate
 from hydra.errors import InstantiationException
 from hydra.types import ConvertMode, TargetConf
 from lightning_utilities.core.rank_zero import rank_zero_only
-from omegaconf import OmegaConf, SCMode
+from omegaconf import DictConfig, OmegaConf, SCMode
 from omegaconf._utils import is_structured_config
 from rich import print
 from rich.panel import Panel
@@ -30,6 +30,12 @@ Function to be used for printing function calls.
 CATCH_EXCEPTION = True
 
 
+def is_instantiable(config: Union[DictConfig, Any]) -> bool:
+    if OmegaConf.is_dict(config):
+        return "_target_" in config
+    return False
+
+
 def _resolve_callable_name(f: Callable[..., Any]) -> str:
     # Get the module name
     module_name = f.__module__

fusion_bench/utils/json.py

@@ -1,3 +1,33 @@
+import json
+from pathlib import Path
+from typing import Any, Union
+
+
+def save_to_json(obj, path: Union[str, Path]):
+    """
+    save an object to a json file
+
+    Args:
+        obj (Any): the object to save
+        path (Union[str, Path]): the path to save the object
+    """
+    with open(path, "w") as f:
+        json.dump(obj, f)
+
+
+def load_from_json(path: Union[str, Path]) -> Union[dict, list]:
+    """load an object from a json file
+
+    Args:
+        path (Union[str, Path]): the path to load the object
+
+    Returns:
+        dict: the loaded object
+    """
+    with open(path, "r") as f:
+        return json.load(f)
+
+
 def _is_list_of_dict(obj) -> bool:
     if not isinstance(obj, list):
         return False

fusion_bench/utils/parameters.py

@@ -14,6 +14,7 @@ __all__ = [
     "get_parameter_statistics",
     "state_dict_to_vector",
     "vector_to_state_dict",
+    "trainable_state_dict",
 ]
 
 # Model conversion utils
@@ -44,33 +45,43 @@ def trainable_state_dict(
 
 
 def state_dict_to_vector(
-    state_dict: StateDictType,
+    state_dict: Union[StateDictType, nn.Module],
     remove_keys: Optional[List[str]] = None,
 ):
     """
     Convert a state dictionary to a vector.
 
     Args:
-        state_dict (dict): The state dictionary to convert.
+        state_dict (Union[dict[str, torch.Tensor], nn.Module]): The state dictionary to convert.
         remove_keys (list, optional): List of keys to remove from the state dictionary. Defaults to [].
 
     Returns:
         torch.Tensor: The converted vector.
     """
     remove_keys = remove_keys if remove_keys is not None else []
-    shared_state_dict = copy.deepcopy(state_dict)
+
+    if isinstance(state_dict, nn.Module):
+        shared_state_dict = state_dict.state_dict()
+    else:
+        shared_state_dict = copy.copy(state_dict)
+
+    # remove the keys to be removed
     for key in remove_keys:
         if key in shared_state_dict:
             del shared_state_dict[key]
+
+    # sort the reference dict
     sorted_shared_state_dict = OrderedDict(sorted(shared_state_dict.items()))
-    return nn.utils.parameters_to_vector(
+
+    vector = nn.utils.parameters_to_vector(
         [value.reshape(-1) for key, value in sorted_shared_state_dict.items()]
     )
+    return vector
 
 
 def vector_to_state_dict(
     vector: torch.Tensor,
-    state_dict: StateDictType,
+    state_dict: Union[StateDictType, nn.Module],
     remove_keys: Optional[List[str]] = None,
 ):
     """
@@ -78,18 +89,27 @@ def vector_to_state_dict(
 
     Args:
         vector (torch.Tensor): The vector to convert.
-        state_dict (dict): The reference state dictionary to define the order of the vector.
+        state_dict (Union[dict[str, torch.Tensor], nn.Module]): The reference state dictionary to define the order of the vector.
         remove_keys (list, optional): List of keys to remove from the reference state dictionary. Defaults to [].
 
     Returns:
         dict: The converted state dictionary.
     """
     remove_keys = remove_keys if remove_keys is not None else []
+
     # create a reference dict to define the order of the vector
-    reference_dict = copy.deepcopy(state_dict)
+    if isinstance(state_dict, nn.Module):
+        reference_dict = state_dict.state_dict()
+    else:
+        # shallow copy the state_dict
+        reference_dict = copy.copy(state_dict)
+
+    # remove the keys to be removed
     for key in remove_keys:
         if key in reference_dict:
             del reference_dict[key]
+
+    # sort the reference dict
     sorted_reference_dict = OrderedDict(sorted(reference_dict.items()))
 
     # create a shared state dict using the reference dict

fusion_bench/utils/path.py

@@ -1,7 +1,22 @@
 import os
+from typing import List
 
 
 def path_is_dir_and_not_empty(path: str):
     if path is None:
         return False
     return os.path.isdir(path) and len(os.listdir(path)) > 0
+
+
+def listdir_fullpath(dir: str) -> List[str]:
+    """list directory `dir`, return fullpaths
+
+    Args:
+        dir (str): directory name
+
+    Returns:
+        List[str]: a list of fullpaths
+    """
+    assert os.path.isdir(dir), "Argument 'dir' must be a Directory"
+    names = os.listdir(dir)
+    return [os.path.join(dir, name) for name in names]