Remove the need to iterate the model parameters for freezing and speedup (#176)

Summary:
Pull Request resolved: #176

in swav, there is a parameter that's frozen for 313 iterations. the previous logic iterated through all model named_parameters and checked if the parameter is frozen + iterations <313. This is inefficient as the model grows bigger in size and iterating through all parameters for the check will be slow.

alternately, now we check whats the max frozen iteration to which a parameter is frozen and we only iterate named_parameters() of the model if the current iteration <= max frozen iteration

inspired by blefaudeux

Reviewed By: min-xu-ai

Differential Revision: D26276803

fbshipit-source-id: 6d91658273e5ea7fb56ea2c796e0707e90c801b7

Author: prigoyal

vissl/hooks/state_update_hooks.py

@@ -260,6 +260,18 @@ def on_backward(self, task: "tasks.ClassyTask") -> None:
         map_params_to_iters = {}
         for to_map in task.config.MODEL.TEMP_FROZEN_PARAMS_ITER_MAP:
             map_params_to_iters[to_map[0]] = to_map[1]
+
+        # get the maximum iterations until which the params are frozen.
+        # if the iterations are past the maximum iterations freezing any
+        # param, we simply return.
+        max_iterations = max(list(map_params_to_iters.values()))
+        if task.iteration >= max_iterations:
+            if task.iteration == max_iterations:
+                logging.info(
+                    f"No parameters grad removed from now on: {task.iteration}"
+                )
+            return
+
         for name, p in task.model.named_parameters():
             if (
                 name in map_params_to_iters