def compute_survival_rate(layer_id, n_layer,
                          survival_rate_last_layer, survival_schedule):
  """Computer survival rate for stochastic depth."""
  if survival_schedule == 'exp':
    # alpha^(n_layer - 1) = survival_rate_last_layer
    # alpha = survival_rate_last_layer^(1/(n_layer-1))
    # survival_rate = alpha^layer_id
    survival_rate = survival_rate_last_layer ** (
        float(layer_id) / (n_layer - 1))
  elif survival_schedule == 'linear':
    survival_rate = 1.0 - float(layer_id) / n_layer * (
        1 - survival_rate_last_layer)
  else:
    raise ValueError('Unknown survival_schedule %s' % survival_schedule)
  return survival_rate