(self, batch, rng, step=0, learning_rate=None)
return stats[0]['loss'], joint_stats
Updates layers weights/state and optimizers slots by running one step.

    Args:
      batch: Batch of data to use for optimization.
      rng: Random number generator to use for running this step.
      step: Which step of the training are we running.
      learning_rate: Learning rate to use instead of the default one.

    Returns:
      Tuple (loss, stats) with new values from one step
      of training, where stats are all optimizer statistics.
Updates layers weights/state and optimizers slots by running one step.
Updates layers weights / state and optimizers slots by running one step .
def one_step(self, batch, rng, step=0, learning_rate=None):
    """Updates layers weights/state and optimizers slots by running one step.

    Args:
      batch: Batch of data to use for optimization.
      rng: Random number generator to use for running this step.
      step: Which step of the training are we running.
      learning_rate: Learning rate to use instead of the default one.

    Returns:
      Tuple (loss, stats) with new values from one step
      of training, where stats are all optimizer statistics.
    """
    # Update the learning rate if needed.
    if learning_rate is not None:
      self._replicated_loss_opt_params['learning_rate'] = self._replicate_cpu(
          learning_rate)
      for (std_op, rev_ops) in self._replicated_opt_params:
        std_op['learning_rate'] = self._replicate_cpu(learning_rate)
        for op in rev_ops:
          op['learning_rate'] = self._replicate_cpu(learning_rate)

    # Batch needs to be split across the local devices -- the difference
    # between _for_n_devices and _reshape_by_device is that the latter splits
    # the batch dim to batch // n_devices, vs _for_n_devices
    # broadcasts/replicates to n_devices dimension.
    step_int = step
    if self._n_devices > 1:
      batch = tl.reshape_by_device(batch, self._n_devices, pure_np=True)
      step = np.repeat(step, self._n_devices)

    # Create separate rng for each device and layer.
    if self._n_devices == 1:
      rngs = fastmath.random.split(rng, self._n_layers)
    else:
      # JIT the function and run it on CPU to avoid memory fragmentation.
      rngs = self._jit_per_device_rngs(tl.on_cpu(rng))
    # Group rngs by layer blocks.
    rng_blocks, rng_i = [], 0
    for _, rev_layers in self._blocks:
      l = len(rev_layers)
      rng_blocks.append((rngs[rng_i], rngs[rng_i + 1: rng_i + l + 1]))
      rng_i += l + 1

    # Run the layers forward upto the loss layer.
    if self._do_free:
      self._free_accelerators()
    process = psutil.Process(os.getpid())
    if isinstance(batch, (list, tuple)):
      batch_shapes = [x.shape for x in batch]
    else:
      batch_shapes = batch.shape
    logging.info('running step %d on shapes %s', step_int, str(batch_shapes))
    if step_int % self._n_steps_per_log == 1:
      logging.info('run fwd: cpu memory use (MB): %.2f',
                   process.memory_info().rss / float(1024 * 1024))

    stack = batch
    block_inputs_states = []
    for i, (std_layer, rev_layers) in enumerate(self._blocks):
      acc_std_layer_fn, acc_rev_layer_fns = self._accelerated_layer_fns[i]
      std_rng, rev_rngs = rng_blocks[i]
      # Run the standard layer.
      stack, std_inputs, std_state = self._run_forward_standard(
          stack, std_layer, acc_std_layer_fn, std_rng, step_int)

      # Run the reversible layers and collect old and new states.
      stack, rev_old_states, rev_new_states = self._run_forward_reversible(
          stack, rev_layers, acc_rev_layer_fns, rev_rngs, step_int)
      block_inputs_states.append(tl.on_cpu(
          ((std_inputs, std_state), (rev_old_states, rev_new_states))))

    # Run the loss layer forward and backward with optimizer update.
    if step_int % self._n_steps_per_log == 1:
      logging.info('run loss: cpu memory use (MB): %.2f',
                   process.memory_info().rss / float(1024 * 1024))
    loss_state = self._replicate(self._loss_layer.state)
    loss_inputs = cb.inputs_from_stack(stack, self._loss_layer.n_in)
    loss_stats, grad_stack = self._run_backward_standard(
        None, step, self._loss_layer, loss_inputs,
        loss_state, self._loss_fbo, rngs[-1], self._loss_opt,
        self._replicated_loss_opt_params)
    self._collect_weights(self._loss_layer)
    stats = [tl.on_cpu(loss_stats)]

    # De-fragment memory.
    if self._do_free:
      stack, grad_stack = tl.on_cpu(stack), tl.on_cpu(grad_stack)
      self._free_accelerators()

    # Run the layers backward and run optimizer updates.
    if step_int % self._n_steps_per_log == 1:
      logging.info('run bwd: cpu memory use (MB): %.2f',
                   process.memory_info().rss / float(1024 * 1024))
    for i in range(len(self._blocks) - 1, -1, -1):
      std_layer, rev_layers = self._blocks[i]
      (std_inputs, std_state), (rev_old_states,
                                rev_new_states) = block_inputs_states[i]
      std_fbo, rev_fbos = self._fbos[i]
      std_opt, rev_opts = self._optimizers[i]
      std_rng, rev_rngs = rng_blocks[i]
      repl_std_opt_params, repl_rev_opts_params = self._replicated_opt_params[i]

      # Run reversible layers backward with optimizer update.
      stack, grad_stack, new_stats = self._run_backward_reversible(
          stack, grad_stack, step, rev_layers, rev_fbos, rev_old_states,
          rev_new_states, rev_rngs, rev_opts, repl_rev_opts_params)
      stats.extend(tl.on_cpu(new_stats))

      # Run the standard layer forward-and-backward pass and optimizer update.
      std_layer_stats, grad_stack = self._run_backward_standard(
          grad_stack, step, std_layer, std_inputs, std_state, std_fbo, std_rng,
          std_opt, repl_std_opt_params)
      stack = cb.outputs_onto_stack(  # Put layer inputs on the stack.
          std_inputs, stack, std_layer.n_out)
      stats.append(tl.on_cpu(std_layer_stats))

      # Collect lazily unreplicated layer weights.
      for rev_layer_id in range(self._n_async_layers):
        self._collect_weights(rev_layers[rev_layer_id])
      self._collect_weights(std_layer)

    # Join stats from different optimizers into one.
    joint_stats = {}
    for i, stat in enumerate(reversed(stats)):
      for k, v in stat.items():
        joint_stats[f'layer{i}/' + k] = v
    return stats[0]['loss'], joint_stats
def one_step ( self , batch , rng , step = 0 , learning_rate = None ) : # Update the learning rate if needed. if learning_rate is not None : self . _replicated_loss_opt_params [ 'learning_rate' ] = self . _replicate_cpu ( learning_rate ) for ( std_op , rev_ops ) in self . _replicated_opt_params : std_op [ 'learning_rate' ] = self . _replicate_cpu ( learning_rate ) for op in rev_ops : op [ 'learning_rate' ] = self . _replicate_cpu ( learning_rate ) # Batch needs to be split across the local devices -- the difference # between _for_n_devices and _reshape_by_device is that the latter splits # the batch dim to batch // n_devices, vs _for_n_devices # broadcasts/replicates to n_devices dimension. step_int = step if self . _n_devices > 1 : batch = tl . reshape_by_device ( batch , self . _n_devices , pure_np = True ) step = np . repeat ( step , self . _n_devices ) # Create separate rng for each device and layer. if self . _n_devices == 1 : rngs = fastmath . random . split ( rng , self . _n_layers ) else : # JIT the function and run it on CPU to avoid memory fragmentation. rngs = self . _jit_per_device_rngs ( tl . on_cpu ( rng ) ) # Group rngs by layer blocks. rng_blocks , rng_i = [ ] , 0 for _ , rev_layers in self . _blocks : l = len ( rev_layers ) rng_blocks . append ( ( rngs [ rng_i ] , rngs [ rng_i + 1 : rng_i + l + 1 ] ) ) rng_i += l + 1 # Run the layers forward upto the loss layer. if self . _do_free : self . _free_accelerators ( ) process = psutil . Process ( os . getpid ( ) ) if isinstance ( batch , ( list , tuple ) ) : batch_shapes = [ x . shape for x in batch ] else : batch_shapes = batch . shape logging . info ( 'running step %d on shapes %s' , step_int , str ( batch_shapes ) ) if step_int % self . _n_steps_per_log == 1 : logging . info ( 'run fwd: cpu memory use (MB): %.2f' , process . memory_info ( ) . rss / float ( 1024 * 1024 ) ) stack = batch block_inputs_states = [ ] for i , ( std_layer , rev_layers ) in enumerate ( self . _blocks ) : acc_std_layer_fn , acc_rev_layer_fns = self . _accelerated_layer_fns [ i ] std_rng , rev_rngs = rng_blocks [ i ] # Run the standard layer. stack , std_inputs , std_state = self . _run_forward_standard ( stack , std_layer , acc_std_layer_fn , std_rng , step_int ) # Run the reversible layers and collect old and new states. stack , rev_old_states , rev_new_states = self . _run_forward_reversible ( stack , rev_layers , acc_rev_layer_fns , rev_rngs , step_int ) block_inputs_states . append ( tl . on_cpu ( ( ( std_inputs , std_state ) , ( rev_old_states , rev_new_states ) ) ) ) # Run the loss layer forward and backward with optimizer update. if step_int % self . _n_steps_per_log == 1 : logging . info ( 'run loss: cpu memory use (MB): %.2f' , process . memory_info ( ) . rss / float ( 1024 * 1024 ) ) loss_state = self . _replicate ( self . _loss_layer . state ) loss_inputs = cb . inputs_from_stack ( stack , self . _loss_layer . n_in ) loss_stats , grad_stack = self . _run_backward_standard ( None , step , self . _loss_layer , loss_inputs , loss_state , self . _loss_fbo , rngs [ - 1 ] , self . _loss_opt , self . _replicated_loss_opt_params ) self . _collect_weights ( self . _loss_layer ) stats = [ tl . on_cpu ( loss_stats ) ] # De-fragment memory. if self . _do_free : stack , grad_stack = tl . on_cpu ( stack ) , tl . on_cpu ( grad_stack ) self . _free_accelerators ( ) # Run the layers backward and run optimizer updates. if step_int % self . _n_steps_per_log == 1 : logging . info ( 'run bwd: cpu memory use (MB): %.2f' , process . memory_info ( ) . rss / float ( 1024 * 1024 ) ) for i in range ( len ( self . _blocks ) - 1 , - 1 , - 1 ) : std_layer , rev_layers = self . _blocks [ i ] ( std_inputs , std_state ) , ( rev_old_states , rev_new_states ) = block_inputs_states [ i ] std_fbo , rev_fbos = self . _fbos [ i ] std_opt , rev_opts = self . _optimizers [ i ] std_rng , rev_rngs = rng_blocks [ i ] repl_std_opt_params , repl_rev_opts_params = self . _replicated_opt_params [ i ] # Run reversible layers backward with optimizer update. stack , grad_stack , new_stats = self . _run_backward_reversible ( stack , grad_stack , step , rev_layers , rev_fbos , rev_old_states , rev_new_states , rev_rngs , rev_opts , repl_rev_opts_params ) stats . extend ( tl . on_cpu ( new_stats ) ) # Run the standard layer forward-and-backward pass and optimizer update. std_layer_stats , grad_stack = self . _run_backward_standard ( grad_stack , step , std_layer , std_inputs , std_state , std_fbo , std_rng , std_opt , repl_std_opt_params ) stack = cb . outputs_onto_stack ( # Put layer inputs on the stack. std_inputs , stack , std_layer . n_out ) stats . append ( tl . on_cpu ( std_layer_stats ) ) # Collect lazily unreplicated layer weights. for rev_layer_id in range ( self . _n_async_layers ) : self . _collect_weights ( rev_layers [ rev_layer_id ] ) self . _collect_weights ( std_layer ) # Join stats from different optimizers into one. joint_stats = { } for i , stat in enumerate ( reversed ( stats ) ) : for k , v in stat . items ( ) : joint_stats [ f'layer{i}/' + k ] = v return stats [ 0 ] [ 'loss' ] , joint_stats