burn_core/module/param/
group.rs

1use crate::module::{Module, ModuleVisitor, Param};
2
3use alloc::string::String;
4#[cfg(target_has_atomic = "ptr")]
5use alloc::sync::Arc;
6use alloc::vec;
7use alloc::vec::Vec;
8#[cfg(not(target_has_atomic = "ptr"))]
9use portable_atomic_util::Arc;
10#[cfg(feature = "std")]
11use regex::Regex;
12
13use burn_std::id::ParamId;
14use burn_tensor::{Bool, Int, Tensor};
15
16/// Errors tied to [ParamGroup]'s.
17#[derive(Debug)]
18pub enum ParamGroupError {
19    /// Use of an invalid Regex pattern.
20    InvalidPatternError(String),
21}
22
23#[derive(Default)]
24struct ParamIdCollector {
25    ids: Vec<ParamId>,
26}
27
28impl ParamIdCollector {
29    pub fn ids(&self) -> Vec<ParamId> {
30        self.ids.clone()
31    }
32}
33
34impl ModuleVisitor for ParamIdCollector {
35    fn visit_float<const D: usize>(&mut self, param: &Param<Tensor<D>>) {
36        self.ids.push(param.id);
37    }
38
39    fn visit_int<const D: usize>(&mut self, param: &Param<Tensor<D, Int>>) {
40        self.ids.push(param.id);
41    }
42
43    fn visit_bool<const D: usize>(&mut self, param: &Param<Tensor<D, Bool>>) {
44        self.ids.push(param.id);
45    }
46}
47
48/// A way to represent a group of parameter for a Burn module.
49#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
50pub struct ParamGroup {
51    matcher: ParamGroupMatcher,
52    excludes: Option<ParamGroupMatcher>,
53}
54
55impl ParamGroup {
56    /// Evaluates whether a given parameter ID and its module path match this group.
57    pub fn matches(&self, id: &ParamId, path: Option<&str>) -> bool {
58        let matched = self.matcher.matches(id, path);
59
60        let excluded = if let Some(exclude_matcher) = &self.excludes {
61            exclude_matcher.matches(id, path)
62        } else {
63            false
64        };
65
66        matched && !excluded
67    }
68
69    /// Returns a parameter group with all of the module's parameters.
70    pub fn ids_from_module<M: Module>(module: M) -> Self {
71        let mut collector = ParamIdCollector::default();
72        module.visit(&mut collector);
73        Self {
74            matcher: ParamGroupMatcher::from_ids(collector.ids()),
75            excludes: None,
76        }
77    }
78
79    /// Matches parameters by exact text path (e.g., "model.backbone.linear.weight")
80    pub fn from_path(path: impl Into<String>) -> Self {
81        ParamGroup::from_paths(vec![path])
82    }
83
84    /// Matches parameters by exact text paths (e.g., "model.backbone.linear.weight", etc.)
85    pub fn from_paths(paths: Vec<impl Into<String>>) -> Self {
86        Self {
87            matcher: ParamGroupMatcher::Path(Arc::new(PathMatcher::Exact(
88                paths.into_iter().map(|p| p.into()).collect(),
89            ))),
90            excludes: None,
91        }
92    }
93
94    /// Matches parameters that include the predicate in their paths (e.g., "backbone")
95    pub fn from_predicate(path: impl Into<String>) -> Self {
96        ParamGroup::from_predicates(vec![path])
97    }
98
99    /// Matches parameters that include all the predicates in their path (AND logic).
100    /// (e.g., parameter path contains "backbone" and "linear")
101    pub fn from_predicates(paths: Vec<impl Into<String>>) -> Self {
102        Self {
103            matcher: ParamGroupMatcher::Path(Arc::new(PathMatcher::Include(
104                paths.into_iter().map(|p| p.into()).collect(),
105            ))),
106            excludes: None,
107        }
108    }
109
110    /// Matches parameters that include any of the predicates in their path (OR logic).
111    /// (e.g., parameter path contains "backbone" or "linear")
112    pub fn from_any_predicates(paths: Vec<impl Into<String>>) -> Self {
113        let mut matchers: Vec<ParamGroupMatcher> = paths
114            .into_iter()
115            .map(|p| ParamGroupMatcher::Path(Arc::new(PathMatcher::Include(vec![p.into()]))))
116            .collect();
117        let mut main_matcher = if let Some(value) = matchers.pop() {
118            value
119        } else {
120            return Self {
121                matcher: ParamGroupMatcher::Path(Arc::new(PathMatcher::Include(vec![]))),
122                excludes: None,
123            };
124        };
125
126        matchers
127            .iter()
128            .for_each(|m| main_matcher = main_matcher.clone().fuse(m));
129
130        Self {
131            matcher: main_matcher,
132            excludes: None,
133        }
134    }
135
136    #[cfg(feature = "std")]
137    /// Matches parameters by regex pattern (e.g., "^model\.layer\.\d+$")
138    ///
139    /// # Errors
140    /// Returns a [ParamGroupError::InvalidPatternError] if the string cannot be compiled into a valid regex.
141    pub fn from_regex<S: AsRef<str>>(pattern: S) -> Result<Self, ParamGroupError> {
142        ParamGroup::from_regexes(vec![pattern])
143    }
144
145    #[cfg(feature = "std")]
146    /// Matches parameters for all the regex patterns (AND logic).
147    /// (e.g., "^encoder\.layer\.\d+", and "bias$" )
148    ///
149    /// # Errors
150    /// Returns a [ParamGroupError::InvalidPatternError] if the strings cannot be compiled into a valid regex.
151    pub fn from_regexes<S: AsRef<str>>(patterns: Vec<S>) -> Result<Self, ParamGroupError> {
152        let mut new_patterns = vec![];
153        for pattern in patterns {
154            match Regex::new(pattern.as_ref()) {
155                Ok(re) => new_patterns.push(re),
156                Err(e) => {
157                    return Err(ParamGroupError::InvalidPatternError(format!(
158                        "Invalid regex pattern: {e}"
159                    )));
160                }
161            }
162        }
163        Ok(Self {
164            matcher: ParamGroupMatcher::Path(Arc::new(PathMatcher::Regex(new_patterns))),
165            excludes: None,
166        })
167    }
168
169    #[cfg(feature = "std")]
170    /// Matches parameters for any the regex patterns (OR logic).
171    /// (e.g., "^encoder\.layer\.\d+$", or "^decoder\.layer\.\d+$" )
172    ///
173    /// # Errors
174    /// Returns a [ParamGroupError::InvalidPatternError] if the strings cannot be compiled into a valid regex.
175    pub fn from_any_regexes<S: AsRef<str>>(patterns: Vec<S>) -> Result<Self, ParamGroupError> {
176        let mut matchers = vec![];
177        for pattern in patterns {
178            match Regex::new(pattern.as_ref()) {
179                Ok(re) => {
180                    matchers.push(ParamGroupMatcher::Path(Arc::new(PathMatcher::Regex(vec![
181                        re,
182                    ]))))
183                }
184                Err(e) => {
185                    return Err(ParamGroupError::InvalidPatternError(format!(
186                        "Invalid regex pattern: {e}"
187                    )));
188                }
189            }
190        }
191
192        let mut main_matcher = if let Some(value) = matchers.pop() {
193            value
194        } else {
195            return Ok(Self {
196                matcher: ParamGroupMatcher::Path(Arc::new(PathMatcher::Include(vec![]))),
197                excludes: None,
198            });
199        };
200
201        matchers
202            .iter()
203            .for_each(|m| main_matcher = main_matcher.clone().fuse(m));
204
205        Ok(Self {
206            matcher: main_matcher,
207            excludes: None,
208        })
209    }
210
211    /// Matches any parameter.
212    pub fn all() -> Self {
213        Self {
214            matcher: ParamGroupMatcher::All,
215            excludes: None,
216        }
217    }
218
219    /// Matches a specific slice of predefined parameter IDs.
220    pub fn from_ids(ids: Vec<ParamId>) -> Self {
221        Self {
222            matcher: ParamGroupMatcher::Explicit(Arc::new(ids)),
223            excludes: None,
224        }
225    }
226
227    /// Fuse two parameter group.
228    pub fn fuse(self, other: &Self) -> Self {
229        Self {
230            matcher: self.matcher.fuse(&other.matcher),
231            excludes: None,
232        }
233    }
234
235    /// Exclude the given group from the current group
236    pub fn exclude(mut self, group: Self) -> Self {
237        self.excludes = match &self.excludes {
238            Some(excluded) => Some(excluded.clone().fuse(&group.matcher)),
239            None => Some(group.matcher.clone()),
240        };
241        self
242    }
243}
244
245mod arc_serde {
246    use serde::{Deserialize, Deserializer, Serialize, Serializer};
247
248    use super::*;
249
250    pub fn serialize<S, T>(val: &Arc<T>, serializer: S) -> Result<S::Ok, S::Error>
251    where
252        S: Serializer,
253        T: Serialize,
254    {
255        val.as_ref().serialize(serializer)
256    }
257
258    pub fn deserialize<'de, D, T>(deserializer: D) -> Result<Arc<T>, D::Error>
259    where
260        D: Deserializer<'de>,
261        T: Deserialize<'de>,
262    {
263        let v = T::deserialize(deserializer)?;
264        Ok(Arc::new(v))
265    }
266}
267
268#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
269enum ParamGroupMatcher {
270    All,
271    #[serde(with = "arc_serde")]
272    Explicit(Arc<Vec<ParamId>>),
273    #[serde(with = "arc_serde")]
274    Path(Arc<PathMatcher>),
275    #[serde(with = "arc_serde")]
276    Combined(Arc<Vec<Self>>),
277}
278
279impl ParamGroupMatcher {
280    pub fn from_ids(ids: Vec<ParamId>) -> Self {
281        Self::Explicit(Arc::new(ids))
282    }
283
284    pub(crate) fn matches(&self, id: &ParamId, path: Option<&str>) -> bool {
285        match self {
286            Self::All => true,
287            Self::Explicit(ids) => ids.contains(id),
288            Self::Path(matcher) => path.is_some_and(|p| matcher.matches(p)),
289            Self::Combined(matchers) => matchers.iter().any(|m| m.matches(id, path)),
290        }
291    }
292
293    fn push_combined(self, other: Self) -> Self {
294        match self {
295            ParamGroupMatcher::Combined(param_group_matchers) => match other.clone() {
296                ParamGroupMatcher::All => Self::All,
297                ParamGroupMatcher::Explicit(_) => {
298                    let mut matchers = (*param_group_matchers).clone();
299                    matchers.push(other);
300                    Self::Combined(Arc::new(matchers))
301                }
302                ParamGroupMatcher::Path(_) => {
303                    let mut matchers = (*param_group_matchers).clone();
304                    matchers.push(other);
305                    Self::Combined(Arc::new(matchers))
306                }
307                ParamGroupMatcher::Combined(other_matchers) => {
308                    let mut matchers = (*param_group_matchers).clone();
309                    matchers.append(&mut (*other_matchers).clone());
310                    Self::Combined(Arc::new(matchers))
311                }
312            },
313            _ => panic!(
314                "`push_combined` should only be called on a ParamGroupMatcher::Combined variant."
315            ),
316        }
317    }
318
319    pub(crate) fn fuse(self, other: &Self) -> Self {
320        match (self.clone(), other.clone()) {
321            (ParamGroupMatcher::All, _) => Self::All,
322            (_, ParamGroupMatcher::All) => Self::All,
323            (ParamGroupMatcher::Explicit(_), ParamGroupMatcher::Combined(_)) => {
324                other.clone().push_combined(self)
325            }
326            (ParamGroupMatcher::Path(_), ParamGroupMatcher::Combined(_)) => {
327                other.clone().push_combined(self)
328            }
329            (ParamGroupMatcher::Combined(_), ParamGroupMatcher::Explicit(_)) => {
330                self.push_combined(other.clone())
331            }
332            (ParamGroupMatcher::Combined(_), ParamGroupMatcher::Path(_)) => {
333                self.push_combined(other.clone())
334            }
335            (ParamGroupMatcher::Combined(_), ParamGroupMatcher::Combined(_)) => {
336                self.push_combined(other.clone())
337            }
338            _ => ParamGroupMatcher::Combined(Arc::new(vec![self, other.clone()])),
339        }
340    }
341}
342
343#[cfg(feature = "std")]
344mod regex_serde {
345    use serde::{Deserialize, Deserializer, Serialize, Serializer};
346
347    use super::*;
348
349    pub fn serialize<S>(regexes: &[Regex], serializer: S) -> Result<S::Ok, S::Error>
350    where
351        S: Serializer,
352    {
353        let v: Vec<&str> = regexes.iter().map(|r| r.as_str()).collect();
354        v.serialize(serializer)
355    }
356
357    pub fn deserialize<'de, D>(deserializer: D) -> Result<Vec<Regex>, D::Error>
358    where
359        D: Deserializer<'de>,
360    {
361        let strings: Vec<String> = Vec::deserialize(deserializer)?;
362        strings
363            .into_iter()
364            .map(|s| Regex::new(&s).map_err(serde::de::Error::custom))
365            .collect()
366    }
367}
368
369#[derive(Clone, serde::Serialize, serde::Deserialize)]
370enum PathMatcher {
371    Exact(Vec<String>),
372    #[cfg(feature = "std")]
373    #[serde(with = "regex_serde")]
374    Regex(Vec<Regex>),
375    Include(Vec<String>),
376}
377
378impl PathMatcher {
379    pub(crate) fn matches(&self, path: &str) -> bool {
380        match self {
381            PathMatcher::Exact(paths) => paths.iter().any(|p| p == path),
382            #[cfg(feature = "std")]
383            PathMatcher::Regex(patterns) => patterns.iter().all(|r| r.is_match(path)),
384            PathMatcher::Include(includes) => includes.iter().all(|inc| path.contains(inc)),
385        }
386    }
387}
388
389impl alloc::fmt::Debug for PathMatcher {
390    fn fmt(&self, f: &mut alloc::fmt::Formatter<'_>) -> alloc::fmt::Result {
391        match self {
392            Self::Exact(arg0) => f.debug_tuple("Exact").field(arg0).finish(),
393            #[cfg(feature = "std")]
394            Self::Regex(arg0) => f.debug_tuple("Regex").field(arg0).finish(),
395            Self::Include(arg0) => f.debug_tuple("Include").field(arg0).finish(),
396        }
397    }
398}
399
400#[cfg(test)]
401mod tests {
402    use super::*;
403    use crate::test_utils::SimpleLinear;
404
405    #[test]
406    fn all_matches_any_parameter() {
407        let group = ParamGroup::all();
408        let id = ParamId::new();
409
410        assert!(group.matches(&id, None));
411    }
412
413    #[test]
414    fn explicit_matches_only_selected_ids() {
415        let id = ParamId::new();
416        let other_id = ParamId::new();
417        let group = ParamGroup::from_ids(vec![id.clone()]);
418
419        assert!(group.matches(&id, None));
420        assert!(!group.matches(&other_id, None));
421    }
422
423    #[test]
424    fn path_matcher_requires_path_and_matches_exactly() {
425        let group = ParamGroup::from_path("model.backbone.weight");
426        let id = ParamId::new();
427
428        assert!(group.matches(&id, Some("model.backbone.weight")));
429        assert!(!group.matches(&id, Some("model.backbone.bias")));
430        assert!(!group.matches(&id, None));
431    }
432
433    #[test]
434    fn predicate_matcher_matches_substrings() {
435        let group = ParamGroup::from_predicate("backbone");
436        let id = ParamId::new();
437
438        assert!(group.matches(&id, Some("model.backbone.weight")));
439        assert!(!group.matches(&id, Some("model.other.weight")));
440    }
441
442    #[cfg(feature = "std")]
443    #[test]
444    fn regex_matcher_matches_pattern() {
445        let group = ParamGroup::from_regex(r"^model\.layer\.[0-9]+\.weight$").unwrap();
446        let id = ParamId::new();
447
448        assert!(group.matches(&id, Some("model.layer.3.weight")));
449        assert!(!group.matches(&id, Some("model.layer.weight")));
450    }
451
452    #[test]
453    fn ids_from_module_collects_all_param_ids() {
454        let device = crate::test_device();
455        let module = SimpleLinear::new(4, 8, &device);
456        let weight_id = module.weight.id;
457        let bias_id = module.bias.as_ref().unwrap().id;
458        let group = ParamGroup::ids_from_module(module);
459
460        assert!(group.matches(&weight_id, Some("weight")));
461        assert!(group.matches(&bias_id, Some("bias")));
462    }
463
464    #[test]
465    fn fuse_combines_multiple_groups() {
466        let id = ParamId::new();
467        let group1 = ParamGroup::from_ids(vec![id.clone()]);
468        let group2 = ParamGroup::from_path("model.layer.weight");
469        let fused = group1.fuse(&group2);
470
471        assert!(fused.matches(&id, Some("model.other.bias")));
472        assert!(fused.matches(&ParamId::new(), Some("model.layer.weight")));
473        assert!(!fused.matches(&ParamId::new(), Some("model.layer.bias")));
474    }
475
476    #[test]
477    fn exclude_removes_matching_ids_from_a_group() {
478        let id = ParamId::new();
479        let excluded_id = ParamId::new();
480        let exclude_group = ParamGroup::from_ids(vec![excluded_id.clone()]);
481        let group =
482            ParamGroup::from_ids(vec![id.clone(), excluded_id.clone()]).exclude(exclude_group);
483
484        assert!(group.matches(&id, None));
485        assert!(!group.matches(&excluded_id, None));
486    }
487
488    #[test]
489    fn exclude_filters_path_matches() {
490        let id = ParamId::new();
491        let group = ParamGroup::from_predicate("backbone")
492            .exclude(ParamGroup::from_path("model.backbone.bias"));
493
494        assert!(group.matches(&id, Some("model.backbone.weight")));
495        assert!(!group.matches(&id, Some("model.backbone.bias")));
496    }
497
498    #[test]
499    fn exclude_ignores_excludes_on_the_provided_group() {
500        let id = ParamId::new();
501        let excluded = ParamGroup::from_path("model.backbone.bias")
502            .exclude(ParamGroup::from_path("model.backbone.weight"));
503        let group = ParamGroup::from_path("model.backbone.weight").exclude(excluded);
504
505        assert!(group.matches(&id, Some("model.backbone.weight")));
506        assert!(!group.matches(&id, Some("model.backbone.bias")));
507    }
508
509    #[test]
510    fn from_any_predicates_matches_any_predicate() {
511        let group = ParamGroup::from_any_predicates(vec!["backbone", "encoder"]);
512        let id = ParamId::new();
513
514        assert!(group.matches(&id, Some("model.backbone.weight")));
515        assert!(group.matches(&id, Some("model.encoder.weight")));
516        assert!(!group.matches(&id, Some("model.decoder.weight")));
517    }
518
519    #[cfg(feature = "std")]
520    #[test]
521    fn from_any_regexes_matches_any_pattern() {
522        let group =
523            ParamGroup::from_any_regexes(vec![r"^model\.layer\.[0-9]+\.weight$", r"^model\.bias$"])
524                .unwrap();
525        let id = ParamId::new();
526
527        assert!(group.matches(&id, Some("model.layer.3.weight")));
528        assert!(group.matches(&id, Some("model.bias")));
529        assert!(!group.matches(&id, Some("model.layer.weight")));
530        assert!(!group.matches(&id, Some("model.other.weight")));
531    }
532
533    #[cfg(feature = "std")]
534    #[test]
535    fn from_regexes_with_invalid_pattern() {
536        let result = ParamGroup::from_regex(r"[invalid(");
537        assert!(result.is_err());
538
539        let result = ParamGroup::from_regexes(vec![r"[invalid("]);
540        assert!(result.is_err());
541
542        let result = ParamGroup::from_any_regexes(vec![r"[invalid("]);
543        assert!(result.is_err());
544    }
545}
burn_core/module/param/group.rs

burn_core/module/param/
group.rs
