import contextlib import copy import functools import threading from contextvars import ContextVar from pathlib import Path from typing import ( Any, Callable, Dict, Generic, Iterable, Iterator, List, Optional, Sequence, Set, Tuple, TypeVar, Union, cast, ) import srsly from .backends import CupyOps, NumpyOps, Ops, ParamServer, get_current_ops from .optimizers import Optimizer # noqa: F401 from .shims import Shim from .types import FloatsXd from .util import ( DATA_VALIDATION, convert_recursive, is_xp_array, partial, validate_fwd_input_output, ) InT = TypeVar("InT") OutT = TypeVar("OutT") SelfT = TypeVar("SelfT", bound="Model") context_operators: ContextVar[dict] = ContextVar("context_operators", default={}) def empty_init(model: "Model", *args, **kwargs) -> "Model": return model class Model(Generic[InT, OutT]): """Class for implementing Thinc models and layers.""" global_id: int = 0 global_id_lock: threading.Lock = threading.Lock() _context_operators = context_operators name: str ops: Ops id: int _func: Callable init: Callable _params: ParamServer _dims: Dict[str, Optional[int]] _layers: List["Model"] _shims: List[Shim] _attrs: Dict[str, Any] _has_params: Dict[str, Optional[bool]] # This "locks" the class, so we get an error if you try to assign to # an unexpected variable. __slots__ = [ "name", "id", "ops", "_func", "init", "_params", "_dims", "_attrs", "_refs", "_layers", "_shims", "_has_params", ] def __init__( self, name: str, forward: Callable, *, init: Optional[Callable] = None, dims: Dict[str, Optional[int]] = {}, params: Dict[str, Optional[FloatsXd]] = {}, layers: Sequence["Model"] = [], shims: List[Shim] = [], attrs: Dict[str, Any] = {}, refs: Dict[str, Optional["Model"]] = {}, ops: Optional[Union[NumpyOps, CupyOps]] = None, ): """Initialize a new model.""" self.name = name if init is None: init = partial(empty_init, self) # Assign to callable attrs: https://github.com/python/mypy/issues/2427 setattr(self, "_func", forward) setattr(self, "init", init) self.ops = ops if ops is not None else get_current_ops() self._params = ParamServer() self._dims = dict(dims) self._attrs = dict(attrs) self._refs = dict(refs) self._layers = list(layers) self._shims = list(shims) # Take care to increment the base class here! It needs to be unique # across all models. with Model.global_id_lock: Model.global_id += 1 self.id = Model.global_id self._has_params = {} for name, value in params.items(): self._has_params[name] = None if value is not None: self.set_param(name, value) @property def layers(self) -> List["Model"]: """A list of child layers of the model. You can append to it to add layers but not reassign it. """ return self._layers @property def shims(self) -> List[Shim]: return self._shims @property def attrs(self) -> Dict[str, Any]: """A dict of the model's attrs. You can write to it to update attrs but not reassign it. """ return self._attrs @property def param_names(self) -> Tuple[str, ...]: """Get the names of registered parameter (including unset).""" return tuple(self._has_params.keys()) @property def grad_names(self) -> Tuple[str, ...]: """Get the names of parameters with registered gradients (including unset).""" return tuple([name for name in self.param_names if self.has_grad(name)]) @property def dim_names(self) -> Tuple[str, ...]: """Get the names of registered dimensions (including unset).""" return tuple(self._dims.keys()) @property def ref_names(self) -> Tuple[str, ...]: """Get the names of registered node references (including unset).""" return tuple(self._refs.keys()) @classmethod @contextlib.contextmanager def define_operators(cls, operators: Dict[str, Callable]): """Bind arbitrary binary functions to Python operators, for use in any `Model` instance. Can (and should) be used as a contextmanager. EXAMPLE: with Model.define_operators({">>": chain}): model = Relu(512) >> Relu(512) >> Softmax() """ token = cls._context_operators.set(dict(operators)) yield cls._context_operators.reset(token) def has_dim(self, name: str) -> Optional[bool]: """Check whether the model has a dimension of a given name. If the dimension is registered but the value is unset, returns None. """ if name not in self._dims: return False elif self._dims[name] is not None: return True else: return None def get_dim(self, name: str) -> int: """Retrieve the value of a dimension of the given name.""" if name not in self._dims: raise KeyError(f"Cannot get dimension '{name}' for model '{self.name}'") value = self._dims[name] if value is None: err = f"Cannot get dimension '{name}' for model '{self.name}': value unset" raise ValueError(err) else: return value def set_dim(self, name: str, value: int, *, force: bool = False) -> None: """Set a value for a dimension.""" if name not in self._dims: raise KeyError( f"Cannot set unknown dimension '{name}' for model '{self.name}'." ) old_value = self._dims[name] has_params = any(bool(y) for x, y in self._has_params.items()) invalid_change = (old_value is not None and old_value != value) and ( not force or force and has_params ) if invalid_change: err = f"Attempt to change dimension '{name}' for model '{self.name}' from {old_value} to {value}" raise ValueError(err) self._dims[name] = value def maybe_get_dim(self, name: str) -> Optional[int]: """Retrieve the value of a dimension of the given name, or None.""" return self.get_dim(name) if self.has_dim(name) else None def has_param(self, name: str) -> Optional[bool]: """Check whether the model has a weights parameter of the given name. Returns None if the parameter is registered but currently unset. """ if name not in self._has_params: return False elif self._has_params[name] is not None: return True else: return None def get_param(self, name: str) -> FloatsXd: """Retrieve a weights parameter by name.""" if name not in self._has_params: raise KeyError(f"Unknown param: '{name}' for model '{self.name}'.") if not self._params.has_param(self.id, name): raise KeyError( f"Parameter '{name}' for model '{self.name}' has not been allocated yet." ) return self._params.get_param(self.id, name) def maybe_get_param(self, name: str) -> Optional[FloatsXd]: """Retrieve a weights parameter by name, or None.""" return self.get_param(name) if self.has_param(name) else None def set_param(self, name: str, value: Optional[FloatsXd]) -> None: """Set a weights parameter's value.""" if value is None: self._has_params[name] = None else: self._params.set_param(self.id, name, value) self._has_params[name] = True def has_grad(self, name: str) -> bool: """Check whether the model has a non-zero gradient for a parameter.""" return self._params.has_grad(self.id, name) def get_grad(self, name: str) -> FloatsXd: """Get a gradient from the model.""" return self._params.get_grad(self.id, name) def set_grad(self, name: str, value: FloatsXd) -> None: """Set a gradient value for the model.""" self._params.set_grad(self.id, name, value) def maybe_get_grad(self, name: str) -> Optional[FloatsXd]: """Retrieve a gradient by name, or None.""" return self.get_grad(name) if self.has_grad(name) else None def inc_grad(self, name: str, value: FloatsXd) -> None: """Increment the gradient of a parameter by a value.""" self._params.inc_grad(self.id, name, value) def has_ref(self, name: str) -> Optional[bool]: """Check whether the model has a reference of a given name. If the reference is registered but the value is unset, returns None. """ if name not in self._refs: return False elif self._refs[name] is not None: return True else: return None def get_ref(self, name: str) -> "Model": """Retrieve the value of a reference of the given name.""" if name not in self._refs: raise KeyError(f"Cannot get reference '{name}' for model '{self.name}'.") value = self._refs[name] if value is None: err = f"Cannot get reference '{name}' for model '{self.name}': value unset." raise ValueError(err) else: return value def maybe_get_ref(self, name: str) -> Optional["Model"]: """Retrieve the value of a reference if it exists, or None.""" return self.get_ref(name) if self.has_ref(name) else None def set_ref(self, name: str, value: Optional["Model"]) -> None: """Set a value for a reference.""" if value is None: self._refs[name] = value elif value in self.walk(): self._refs[name] = value else: raise ValueError("Cannot add reference to node not in tree.") def __call__(self, X: InT, is_train: bool) -> Tuple[OutT, Callable]: """Call the model's `forward` function, returning the output and a callback to compute the gradients via backpropagation.""" return self._func(self, X, is_train=is_train) def initialize(self, X: Optional[InT] = None, Y: Optional[OutT] = None) -> "Model": """Finish initialization of the model, optionally providing a batch of example input and output data to perform shape inference.""" if DATA_VALIDATION.get(): validate_fwd_input_output(self.name, self._func, X, Y) if self.init is not None: self.init(self, X=X, Y=Y) return self def begin_update(self, X: InT) -> Tuple[OutT, Callable[[OutT], InT]]: """Run the model over a batch of data, returning the output and a callback to complete the backward pass. A tuple (Y, finish_update), where Y is a batch of output data, and finish_update is a callback that takes the gradient with respect to the output and an optimizer function, and returns the gradient with respect to the input. """ return self._func(self, X, is_train=True) def predict(self, X: InT) -> OutT: """Call the model's `forward` function with `is_train=False`, and return only the output, instead of the `(output, callback)` tuple. """ return self._func(self, X, is_train=False)[0] def finish_update(self, optimizer: Optimizer) -> None: """Update parameters with current gradients. The optimizer is called with each parameter and gradient of the model. """ for node in self.walk(): for shim in node.shims: shim.finish_update(optimizer) for node in self.walk(): for name in node.param_names: if node.has_grad(name): param, grad = optimizer( (node.id, name), node.get_param(name), node.get_grad(name) ) node.set_param(name, param) @contextlib.contextmanager def use_params(self, params: Dict[Tuple[int, str], FloatsXd]): """Context manager to temporarily set the model's parameters to specified values. The params are a dictionary keyed by model IDs, whose values are arrays of weight values. """ backup = {} for name in self.param_names: key = (self.id, name) if key in params: backup[name] = self.get_param(name) self.set_param(name, params[key]) with contextlib.ExitStack() as stack: for layer in self.layers: stack.enter_context(layer.use_params(params)) for shim in self.shims: stack.enter_context(shim.use_params(params)) yield if backup: for name, param in backup.items(): self.set_param(name, param) def walk(self, *, order: str = "bfs") -> Iterable["Model"]: """Iterate out layers of the model. Nodes are returned in breadth-first order by default. Other possible orders are "dfs_pre" (depth-first search in preorder) and "dfs_post" (depth-first search in postorder).""" if order == "bfs": return self._walk_bfs() elif order == "dfs_pre": return self._walk_dfs(post_order=False) elif order == "dfs_post": return self._walk_dfs(post_order=True) else: raise ValueError("Invalid order, must be one of: bfs, dfs_pre, dfs_post") def _walk_bfs(self) -> Iterable["Model"]: """Iterate out layers of the model, breadth-first.""" queue = [self] seen: Set[int] = set() for node in queue: if id(node) in seen: continue seen.add(id(node)) yield node queue.extend(node.layers) def _walk_dfs(self, post_order: bool = False) -> Iterable["Model"]: """Iterate out layers of the model, depth-first.""" seen: Dict[int, Iterator["Model"]] = dict() stack = [self] seen[id(self)] = iter(self.layers) if not post_order: yield self while stack: try: next_child = next(seen[id(stack[-1])]) if not id(next_child) in seen: if not post_order: yield next_child stack.append(next_child) seen[id(next_child)] = iter(next_child.layers) except StopIteration: if post_order: yield stack[-1] stack.pop() def remove_node(self, node: "Model") -> None: """Remove a node from all layers lists, and then update references. References that no longer point to a node within the tree will be set to `None`. For instance, let's say a node has its grandchild as a reference. If the child is removed, the grandchild reference will be left dangling, so will be set to None. """ for child in list(self.walk()): while node in child.layers: child.layers.remove(node) tree = set(self.walk()) for node in tree: for name in node.ref_names: ref = node.get_ref(name) if ref is not None and ref not in tree: node.set_ref(name, None) def replace_callbacks( self, forward: Callable, *, init: Optional[Callable] = None ) -> None: setattr(self, "_func", forward) setattr(self, "init", init) def replace_node(self, old: "Model", new: "Model") -> bool: """Replace a node anywhere it occurs within the model. Returns a boolean indicating whether the replacement was made.""" seen = False # We need to replace nodes in topological order of the transposed graph # to ensure that a node's dependencies are processed before the node. # This is equivalent to a post-order traversal of the original graph. for node in list(self.walk(order="dfs_post")): if node is old: seen = True else: node._layers = [ new if layer is old else layer for layer in node._layers ] for name in node.ref_names: if node.get_ref(name) is old: node.set_ref(name, new) return seen def get_gradients(self) -> Dict[Tuple[int, str], Tuple[FloatsXd, FloatsXd]]: """Get non-zero gradients of the model's parameters, as a dictionary keyed by the parameter ID. The values are (weights, gradients) tuples. """ gradients = {} for node in self.walk(): for name in node.grad_names: param = node.get_param(name) grad = node.get_grad(name) gradients[(node.id, name)] = (param, grad) return gradients def copy(self: SelfT) -> SelfT: """ Create a copy of the model, its attributes, and its parameters. Any child layers will also be deep-copied. The copy will receive a distinct `model.id` value. """ return self._copy() def _copy( self: SelfT, seen: Optional[Dict[int, Union["Model", Shim]]] = None ) -> SelfT: if seen is None: seen = {} params = {} for name in self.param_names: params[name] = self.get_param(name) if self.has_param(name) else None copied_layers: List[Model] = [] for layer in self.layers: if id(layer) in seen: copied_layers.append(cast(Model, seen[id(layer)])) else: copied_layer = layer._copy(seen) seen[id(layer)] = copied_layer copied_layers.append(copied_layer) copied_shims = [] for shim in self.shims: if id(shim) in seen: copied_shims.append(cast(Shim, seen[id(shim)])) else: copied_shim = shim.copy() seen[id(shim)] = copied_shim copied_shims.append(copied_shim) copied: Model[InT, OutT] = Model( self.name, self._func, init=self.init, params=copy.deepcopy(params), dims=copy.deepcopy(self._dims), attrs=copy.deepcopy(self._attrs), layers=copied_layers, shims=copied_shims, ) for name in self.grad_names: copied.set_grad(name, self.get_grad(name).copy()) return cast(SelfT, copied) def to_gpu(self, gpu_id: int) -> None: # pragma: no cover """Transfer the model to a given GPU device.""" import cupy.cuda.device with cupy.cuda.device.Device(gpu_id): self._to_ops(CupyOps()) def to_cpu(self) -> None: # pragma: no cover """Transfer the model to CPU.""" self._to_ops(NumpyOps()) def _to_ops(self, ops: Ops) -> None: # pragma: no cover """Common method for to_cpu/to_gpu.""" for node in self.walk(): node.ops = ops for name in node.param_names: if node.has_param(name): node.set_param(name, ops.asarray_f(node.get_param(name))) if node.has_grad(name): node.set_grad(name, ops.asarray_f(node.get_grad(name))) for shim in node.shims: shim.to_device(ops.device_type, ops.device_id) def to_bytes(self) -> bytes: """Serialize the model to a bytes representation. Models are usually serialized using msgpack, so you should be able to call msgpack.loads() on the data and get back a dictionary with the contents. Serialization should round-trip identically, i.e. the same bytes should result from loading and serializing a model. """ msg = self.to_dict() to_numpy_le = partial(self.ops.to_numpy, byte_order="<") msg = convert_recursive(is_xp_array, to_numpy_le, msg) return srsly.msgpack_dumps(msg) def to_disk(self, path: Union[Path, str]) -> None: """Serialize the model to disk. Most models will serialize to a single file, which should just be the bytes contents of model.to_bytes(). """ path = Path(path) if isinstance(path, str) else path with path.open("wb") as file_: file_.write(self.to_bytes()) def to_dict(self) -> Dict: """Serialize the model to a dict representation. Serialization should round-trip identically, i.e. the same dict should result from loading and serializing a model. """ # We separate out like this to make it easier to read the data in chunks. # The shims might have large weights, while the nodes data will be # small. The attrs are probably not very large, but could be. # The lists are aligned, and refer to the order of self.walk(). msg: Dict[str, List] = {"nodes": [], "attrs": [], "params": [], "shims": []} nodes = list(self.walk()) # Serialize references by their index into the flattened tree. # This is the main reason we can't accept out-of-tree references: # we'd have no way to serialize/deserialize them. node_to_i: Dict[int, Optional[int]] node_to_i = {node.id: i for i, node in enumerate(nodes)} for i, node in enumerate(nodes): refs: Dict[str, Optional[int]] = {} invalid_refs: List[str] = [] for name in node.ref_names: if not node.has_ref(name): refs[name] = None else: ref = node.get_ref(name) if ref.id in node_to_i: refs[name] = node_to_i[ref.id] else: invalid_refs.append(name) if invalid_refs: raise ValueError(f"Cannot get references: {invalid_refs}") dims = {} for dim in node.dim_names: dims[dim] = node.get_dim(dim) if node.has_dim(dim) else None msg["nodes"].append( {"index": i, "name": node.name, "dims": dims, "refs": refs} ) for node in nodes: attrs = {} for name, value in node.attrs.items(): try: attrs[name] = serialize_attr(value, value, name, node) except TypeError: continue msg["attrs"].append(attrs) for node in nodes: msg["shims"].append([shim.to_bytes() for shim in node.shims]) for node in nodes: params: Dict[str, Optional[FloatsXd]] = {} for name in node.param_names: if node.has_param(name): params[name] = cast(Optional[FloatsXd], node.get_param(name)) else: params[name] = None msg["params"].append(params) return msg def from_bytes(self, bytes_data: bytes) -> "Model": """Deserialize the model from a bytes representation. Models are usually serialized using msgpack, so you should be able to call msgpack.loads() on the data and get back a dictionary with the contents. Serialization should round-trip identically, i.e. the same bytes should result from loading and serializing a model. """ msg = srsly.msgpack_loads(bytes_data) msg = convert_recursive(is_xp_array, self.ops.asarray, msg) return self.from_dict(msg) def from_disk(self, path: Union[Path, str]) -> "Model": """Deserialize the model from disk. Most models will serialize to a single file, which should just be the bytes contents of model.to_bytes(). """ path = Path(path) if isinstance(path, str) else path with path.open("rb") as file_: bytes_data = file_.read() return self.from_bytes(bytes_data) def from_dict(self, msg: Dict) -> "Model": if "nodes" not in msg.keys(): # pragma: no cover err = "Trying to read a Model that was created with an incompatible version of Thinc" raise ValueError(err) nodes = list(self.walk()) if len(msg["nodes"]) != len(nodes): raise ValueError("Cannot deserialize model: mismatched structure") for i, node in enumerate(nodes): info = msg["nodes"][i] node.name = info["name"] for dim, value in info["dims"].items(): if value is not None: node.set_dim(dim, value) for ref, ref_index in info["refs"].items(): if ref_index is None: node.set_ref(ref, None) else: node.set_ref(ref, nodes[ref_index]) for attr, value in msg["attrs"][i].items(): default_value = node.attrs.get(attr) loaded_value = deserialize_attr(default_value, value, attr, node) node.attrs[attr] = loaded_value for param_name, value in msg["params"][i].items(): if value is not None: value = node.ops.asarray(value).copy() node.set_param(param_name, value) for i, shim_bytes in enumerate(msg["shims"][i]): node.shims[i].from_bytes(shim_bytes) return self def can_from_disk(self, path: Union[Path, str], *, strict: bool = True) -> bool: """Check whether serialized data on disk is compatible with the model. If 'strict', the function returns False if the model has an attribute already loaded that would be changed. """ path = Path(path) if isinstance(path, str) else path if path.is_dir() or not path.exists(): return False with path.open("rb") as file_: bytes_data = file_.read() return self.can_from_bytes(bytes_data, strict=strict) def can_from_bytes(self, bytes_data: bytes, *, strict: bool = True) -> bool: """Check whether the bytes data is compatible with the model. If 'strict', the function returns False if the model has an attribute already loaded that would be changed. """ try: msg = srsly.msgpack_loads(bytes_data) except ValueError: return False return self.can_from_dict(msg, strict=strict) def can_from_dict(self, msg: Dict, *, strict: bool = True) -> bool: """Check whether a dictionary is compatible with the model. If 'strict', the function returns False if the model has an attribute already loaded that would be changed. """ if "nodes" not in msg.keys(): return False nodes = list(self.walk()) if len(msg["nodes"]) != len(nodes): return False for i, node in enumerate(nodes): info = msg["nodes"][i] if strict and info["name"] != node.name: return False if len(msg["shims"][i]) != len(node.shims): # TODO: The shims should have a check for this too, but # for now we just check if the lengths match. return False for dim, value in info["dims"].items(): has_dim = node.has_dim(dim) if has_dim is False: return False elif has_dim and node.get_dim(dim) != value: return False for param_name, value in msg["params"][i].items(): has_param = node.has_param(param_name) if has_param is False: return False elif has_param and value is not None: param = node.get_param(param_name) if param.shape != value.shape: return False if strict: for attr, value in msg["attrs"][i].items(): if attr in node.attrs: try: serialized = serialize_attr( node.attrs[attr], node.attrs[attr], attr, node ) except TypeError: continue if serialized != value: return False return True def __add__(self, other: Any) -> "Model": """Apply the function bound to the '+' operator.""" if "+" not in self._context_operators.get(): raise TypeError("Undefined operator: +") return self._context_operators.get()["+"](self, other) def __sub__(self, other: Any) -> "Model": """Apply the function bound to the '-' operator.""" if "-" not in self._context_operators.get(): raise TypeError("Undefined operator: -") return self._context_operators.get()["-"](self, other) def __mul__(self, other: Any) -> "Model": """Apply the function bound to the '*' operator.""" if "*" not in self._context_operators.get(): raise TypeError("Undefined operator: *") return self._context_operators.get()["*"](self, other) def __matmul__(self, other: Any) -> "Model": """Apply the function bound to the '@' operator.""" if "@" not in self._context_operators.get(): raise TypeError("Undefined operator: @") return self._context_operators.get()["@"](self, other) def __div__(self, other: Any) -> "Model": # pragma: no cover """Apply the function bound to the '/' operator.""" if "/" not in self._context_operators.get(): raise TypeError("Undefined operator: /") return self._context_operators.get()["/"](self, other) def __truediv__(self, other: Any) -> "Model": """Apply the function bound to the '/' operator.""" if "/" not in self._context_operators.get(): raise TypeError("Undefined operator: /") return self._context_operators.get()["/"](self, other) def __floordiv__(self, other: Any) -> "Model": """Apply the function bound to the '//' operator.""" if "//" not in self._context_operators.get(): raise TypeError("Undefined operator: //") return self._context_operators.get()["//"](self, other) def __mod__(self, other: Any) -> "Model": """Apply the function bound to the '%' operator.""" if "%" not in self._context_operators.get(): raise TypeError("Undefined operator: %") return self._context_operators.get()["%"](self, other) def __pow__(self, other: Any, **kwargs) -> "Model": """Apply the function bound to the '**' operator.""" if "**" not in self._context_operators.get(): raise TypeError("Undefined operator: **") return self._context_operators.get()["**"](self, other) def __lshift__(self, other: Any) -> "Model": """Apply the function bound to the '<<' operator.""" if "<<" not in self._context_operators.get(): raise TypeError("Undefined operator: <<") return self._context_operators.get()["<<"](self, other) def __rshift__(self, other: Any) -> "Model": """Apply the function bound to the '>>' operator.""" if ">>" not in self._context_operators.get(): raise TypeError("Undefined operator: >>") return self._context_operators.get()[">>"](self, other) def __and__(self, other: Any) -> "Model": """Apply the function bound to the '&' operator.""" if "&" not in self._context_operators.get(): raise TypeError("Undefined operator: &") return self._context_operators.get()["&"](self, other) def __xor__(self, other: Any) -> "Model": """Apply the function bound to the '^' operator.""" if "^" not in self._context_operators.get(): raise TypeError("Undefined operator: ^") return self._context_operators.get()["^"](self, other) def __or__(self, other: Any) -> "Model": """Apply the function bound to the '|' operator.""" if "|" not in self._context_operators.get(): raise TypeError("Undefined operator: |") return self._context_operators.get()["|"](self, other) @functools.singledispatch def serialize_attr(_: Any, value: Any, name: str, model: Model) -> bytes: """Serialize an attribute value (defaults to msgpack). You can register custom serializers using the @serialize_attr.register decorator with the type to serialize, e.g.: @serialize_attr.register(MyCustomObject). """ return srsly.msgpack_dumps(value) @functools.singledispatch def deserialize_attr(_: Any, value: Any, name: str, model: Model) -> Any: """Deserialize an attribute value (defaults to msgpack). You can register custom deserializers using the @deserialize_attr.register decorator with the type to deserialize, e.g.: @deserialize_attr.register(MyCustomObject). """ return srsly.msgpack_loads(value) _ModelT = TypeVar("_ModelT", bound=Model) def change_attr_values(model: _ModelT, mapping: Dict[str, Dict[str, Any]]) -> _ModelT: """Walk over the model's nodes, changing the value of attributes using the provided mapping, which maps node names to attr names to attr values. """ for node in model.walk(): if node.name in mapping: attrs = mapping[node.name] for attr, value in attrs.items(): if attr in node.attrs: node.attrs[attr] = value return model def set_dropout_rate(model: _ModelT, drop: float, attrs=["dropout_rate"]) -> _ModelT: """Walk over the model's nodes, setting the dropout rate. You can specify one or more attribute names, by default it looks for ["dropout_rate"]. """ for node in model.walk(): for attr in attrs: if attr in node.attrs: node.attrs[attr] = drop return model def wrap_model_recursive(model: Model, wrapper: Callable[[Model], _ModelT]) -> _ModelT: """Recursively wrap a model and its submodules. The model is updated in-place.""" for node in list(model.walk()): model.replace_node(node, wrapper(node)) return wrapper(model) __all__ = [ "Model", "serialize_attr", "deserialize_attr", "change_attr_values", "set_dropout_rate", "wrap_model_recursive", ]