docstr: Add hook before forward propagation of each layer to block forward computation until the push-pull and
  parameter update is finished. The blocking is implemented using a lock.
docstr_tok: ["Add","hook","before","forward","propagation","of","each","layer","to","block","forward","computation","until","the","push-pull","and","parameter","update","is","finished",".","The","blocking","is","implemented","using","a","lock","."]
code: def _register_forward_hooks(self):
      """Add hook before forward propagation of each layer to block forward computation until the push-pull and
      parameter update is finished. The blocking is implemented using a lock.
      """
      # Recursively find all submodules
      submodules = []
      q = queue.LifoQueue()
      for mod in self._model.children():
          q.put(mod)
      while not q.empty():
          mod = q.get()
          if len(list(mod.children())) == 0:
              submodules.append(mod)
          else:
              for m in mod.children():
                  q.put(m)
  
      def pre_forward_hook(mod, input):
          for p in mod.parameters():
              fp32_p = self._fp16_to_fp32_map[p]
              if fp32_p in self._handles:
                  while True:
                      with self._lock:
                          if fp32_p not in self._handles:
                              break
                          if self._try_to_synchronize(fp32_p):
                              break
                          else:
                              # In order not to waste GPU cycles, we find another handle to synchronize
                              params = list(self._handles.keys())
                              for other_p in params:
                                  if (
                                      other_p.__hash__()
                                      == fp32_p.__hash__()
                                  ):
                                      continue
                                  if self._try_to_synchronize(other_p):
                                      break
                      time.sleep(0.0001)
  
      def after_forward_hook(mod, input, result):
          for p in mod.parameters():
              self._zero_one_grad(p)
  
      # Register pre-hook and hook for each module
      for mod in reversed(submodules):
          mod.register_forward_pre_hook(pre_forward_hook)
          mod.register_forward_hook(after_forward_hook)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from byteps.torch import _DistributedOptimizer
  from byteps.torch.compression import Compression
  from byteps.torch.ops import push_pull_async_inplace
  from byteps.torch.ops import push_pull
  from byteps.torch.ops import poll, synchronize, declare
  from byteps.torch.ops import init, shutdown
  from byteps.torch.ops import size, local_size, rank, local_rank
  import time
  import threading
  import torch
  import collections
label_desc: Use image node_id 4 for calling the _HalfPrecisionDistributedOptimizer obj's underlying member method code with example usage: obj._register_forward_hooks() without return types