add regex search support for param store
from __future__ import absolute_import, division, print_function
+import re
import weakref
import torch
@@ -84,6 +85,8 @@ class ParamStoreDict(object):
:type name: str
:param init_tensor: initial tensor
:type init_tensor: torch.Tensor
+ :param constraint: torch constraint
+ :type constraint: torch.distributions.constraints.Constraint
:returns: parameter
:rtype: torch.Tensor
"""
@@ -120,6 +123,19 @@ class ParamStoreDict(object):
return param
+ def match(self, name):
+ """
+ Get all parameters that match regex. The parameter must exist.
+
+ :param name: regular expression
+ :type name: str
+ :returns: dict with key param name and value torch Tensor
+ """
+ pattern = re.compile(name)
+ params_dict = {key: self.get_param(key) for key in self._params.keys()
+ if pattern.match(key)}
+ return params_dict
+
def param_name(self, p):
"""
Get parameter name from parameter
