make sure parameters of target encoder is never updated by default

lucidrains · Nov 1, 2020 · c914a18 · c914a18
1 parent 7182100
commit c914a18
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/byol_pytorch/byol_pytorch.py b/byol_pytorch/byol_pytorch.py
@@ -34,6 +34,10 @@ def wrapper(self, *args, **kwargs):
 def get_module_device(module):
     return next(module.parameters()).device
 
+def set_requires_grad(model, val):
+    for p in model.parameters():
+        p.requires_grad = val
+
 # loss fn
 
 def loss_fn(x, y):
@@ -181,6 +185,7 @@ def __init__(self, net, image_size, hidden_layer = -2, projection_size = 256, pr
     @singleton('target_encoder')
     def _get_target_encoder(self):
         target_encoder = copy.deepcopy(self.online_encoder)
+        set_requires_grad(target_encoder, False)
         return target_encoder
 
     def reset_moving_average(self):

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'byol-pytorch',
   packages = find_packages(exclude=['examples']),
-  version = '0.3.1',
+  version = '0.3.2',
   license='MIT',
   description = 'Self-supervised contrastive learning made simple',
   author = 'Phil Wang',