(blocks, loss_layer, input_signature, rng)
Initialize reversible blocks and the loss layer and place weights on CPU.

  Args:
    blocks: List of reversible blocks (pairs of layer lists).
    loss_layer: The final loss layer to initialize.
    input_signature: The signature of the input to the blocks.
    rng: Random key used to initialize the layers.
Initialize reversible blocks and the loss layer and place weights on CPU.
Initialize reversible blocks and the loss layer and place weights on CPU .
def init_reversible_blocks(blocks, loss_layer, input_signature, rng):
  """Initialize reversible blocks and the loss layer and place weights on CPU.

  Args:
    blocks: List of reversible blocks (pairs of layer lists).
    loss_layer: The final loss layer to initialize.
    input_signature: The signature of the input to the blocks.
    rng: Random key used to initialize the layers.
  """
  sig_stack = input_signature
  process = psutil.Process(os.getpid())
  mem_use = process.memory_info().rss
  for (std_layers, rev_layers) in blocks:
    rngs = fastmath.random.split(rng, len(std_layers) + len(rev_layers) + 1)
    rng = rngs[0]
    for layer, layer_rng in zip(std_layers + rev_layers, rngs[1:]):
      sig = cb.inputs_from_stack(sig_stack, layer.n_in)
      layer.init(sig, rng=layer_rng)
      layer.weights = tl.on_cpu(layer.weights)  # store weights in cpu memory
      layer.state = tl.on_cpu(layer.state)  # store weights in cpu memory
      logging.info('init: layer %s\nadded cpu memory (MB): %.2f', str(layer),
                   (process.memory_info().rss - mem_use) / float(1024 * 1024))
      mem_use = process.memory_info().rss
      logging.info('init: cpu memory use (MB): %.2f',
                   mem_use / float(1024 * 1024))
      out_sig = layer.output_signature(sig)
      sig_stack = cb.outputs_onto_stack(out_sig, sig_stack, layer.n_in)
  loss_layer.init(cb.inputs_from_stack(sig_stack, loss_layer.n_in), rng=rng)
  loss_layer.weights = tl.on_cpu(loss_layer.weights)
  loss_layer.state = tl.on_cpu(loss_layer.state)
def init_reversible_blocks ( blocks , loss_layer , input_signature , rng ) : sig_stack = input_signature process = psutil . Process ( os . getpid ( ) ) mem_use = process . memory_info ( ) . rss for ( std_layers , rev_layers ) in blocks : rngs = fastmath . random . split ( rng , len ( std_layers ) + len ( rev_layers ) + 1 ) rng = rngs [ 0 ] for layer , layer_rng in zip ( std_layers + rev_layers , rngs [ 1 : ] ) : sig = cb . inputs_from_stack ( sig_stack , layer . n_in ) layer . init ( sig , rng = layer_rng ) layer . weights = tl . on_cpu ( layer . weights ) # store weights in cpu memory layer . state = tl . on_cpu ( layer . state ) # store weights in cpu memory logging . info ( 'init: layer %s\nadded cpu memory (MB): %.2f' , str ( layer ) , ( process . memory_info ( ) . rss - mem_use ) / float ( 1024 * 1024 ) ) mem_use = process . memory_info ( ) . rss logging . info ( 'init: cpu memory use (MB): %.2f' , mem_use / float ( 1024 * 1024 ) ) out_sig = layer . output_signature ( sig ) sig_stack = cb . outputs_onto_stack ( out_sig , sig_stack , layer . n_in ) loss_layer . init ( cb . inputs_from_stack ( sig_stack , loss_layer . n_in ) , rng = rng ) loss_layer . weights = tl . on_cpu ( loss_layer . weights ) loss_layer . state = tl . on_cpu ( loss_layer . state )