Small module broadcast utility.

MillionIntegrals · Apr 7, 2019 · 2547da7 · 2547da7
1 parent c013955
commit 2547da7
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/vel/util/module_util.py b/vel/util/module_util.py
@@ -26,6 +26,16 @@ def apply_leaf(module, f):
         apply_leaf(l, f)
 
 
+def module_apply_broadcast(m, broadcast_fn, args, kwargs):
+    if hasattr(m, broadcast_fn):
+        getattr(m, broadcast_fn)(*args, **kwargs)
+
+
+def module_broadcast(m, broadcast_fn, *args, **kwargs):
+    """ Call given function in all submodules with given parameters """
+    apply_leaf(m, lambda x: module_apply_broadcast(x, broadcast_fn, args, kwargs))
+
+
 def set_train_mode(module):
     # Only fix ones which we don't want to "train"
     if hasattr(module, 'running_mean') and (getattr(module, 'bn_freeze', False) or not getattr(module, 'trainable', True)):