[reformat] ruff reformat lion submission

Author: tfaod

reference_algorithms/paper_baselines/lion/pytorch/submission.py

@@ -18,15 +18,16 @@
 
 # default Lion parameters
 HPARAMS = {
-    "dropout_rate": 0.1,
-    "learning_rate": 2e-4,
-    "one_minus_beta1": 0.05,
-    "beta2": 0.98,
-    "weight_decay": 0.5,
-    "warmup_factor": 0.02
+  'dropout_rate': 0.1,
+  'learning_rate': 2e-4,
+  'one_minus_beta1': 0.05,
+  'beta2': 0.98,
+  'weight_decay': 0.5,
+  'warmup_factor': 0.02,
 }
 HPARAMS = collections.namedtuple('Hyperparameters', HPARAMS.keys())(**HPARAMS)
 
+
 # Modified from https://github.com/google/automl/blob/master/lion/lion_pytorch.py.
 class Lion(Optimizer):
   def __init__(
@@ -90,11 +91,13 @@ def step(self, closure=None):
     return loss
 
 
-def init_optimizer_state(workload: spec.Workload,
-                         model_params: spec.ParameterContainer,
-                         model_state: spec.ModelAuxiliaryState,
-                         hyperparameters: spec.Hyperparameters,
-                         rng: spec.RandomState) -> spec.OptimizerState:
+def init_optimizer_state(
+  workload: spec.Workload,
+  model_params: spec.ParameterContainer,
+  model_state: spec.ModelAuxiliaryState,
+  hyperparameters: spec.Hyperparameters,
+  rng: spec.RandomState,
+) -> spec.OptimizerState:
   """Creates a Lion optimizer and a learning rate schedule."""
   del model_state
   del rng
@@ -103,44 +106,47 @@ def init_optimizer_state(workload: spec.Workload,
   hyperparameters = HPARAMS
 
   optimizer_state = {
-      'optimizer':
-          Lion(
-              model_params.parameters(),
-              lr=HPARAMS.learning_rate,
-              betas=(1.0 - HPARAMS.one_minus_beta1,
-                     HPARAMS.beta2),
-              weight_decay=HPARAMS.weight_decay)
+    'optimizer': Lion(
+      model_params.parameters(),
+      lr=HPARAMS.learning_rate,
+      betas=(1.0 - HPARAMS.one_minus_beta1, HPARAMS.beta2),
+      weight_decay=HPARAMS.weight_decay,
+    )
   }
 
   def pytorch_cosine_warmup(step_hint: int, hyperparameters, optimizer):
     warmup_steps = int(hyperparameters.warmup_factor * step_hint)
     warmup = LinearLR(
-        optimizer, start_factor=1e-10, end_factor=1., total_iters=warmup_steps)
+      optimizer, start_factor=1e-10, end_factor=1.0, total_iters=warmup_steps
+    )
     cosine_steps = max(step_hint - warmup_steps, 1)
     cosine_decay = CosineAnnealingLR(optimizer, T_max=cosine_steps)
     return SequentialLR(
-        optimizer, schedulers=[warmup, cosine_decay], milestones=[warmup_steps])
+      optimizer, schedulers=[warmup, cosine_decay], milestones=[warmup_steps]
+    )
 
   optimizer_state['scheduler'] = pytorch_cosine_warmup(
-      workload.step_hint, HPARAMS, optimizer_state['optimizer'])
+    workload.step_hint, HPARAMS, optimizer_state['optimizer']
+  )
   optimizer_state['hyperparameters'] = hyperparameters
 
   return optimizer_state
 
 
 def update_params(
-    workload: spec.Workload,
-    current_param_container: spec.ParameterContainer,
-    current_params_types: spec.ParameterTypeTree,
-    model_state: spec.ModelAuxiliaryState,
-    hyperparameters: spec.Hyperparameters,
-    batch: Dict[str, spec.Tensor],
-    loss_type: spec.LossType,
-    optimizer_state: spec.OptimizerState,
-    eval_results: List[Tuple[int, float]],
-    global_step: int,
-    rng: spec.RandomState,
-    train_state: Optional[Dict[str, Any]] = None) -> spec.UpdateReturn:
+  workload: spec.Workload,
+  current_param_container: spec.ParameterContainer,
+  current_params_types: spec.ParameterTypeTree,
+  model_state: spec.ModelAuxiliaryState,
+  hyperparameters: spec.Hyperparameters,
+  batch: Dict[str, spec.Tensor],
+  loss_type: spec.LossType,
+  optimizer_state: spec.OptimizerState,
+  eval_results: List[Tuple[int, float]],
+  global_step: int,
+  rng: spec.RandomState,
+  train_state: Optional[Dict[str, Any]] = None,
+) -> spec.UpdateReturn:
   """Return (updated_optimizer_state, updated_params, updated_model_state)."""
   del current_params_types
   del loss_type
@@ -155,26 +161,30 @@ def update_params(
   optimizer_state['optimizer'].zero_grad()
 
   logits_batch, new_model_state = workload.model_fn(
-      params=current_model,
-      augmented_and_preprocessed_input_batch=batch,
-      model_state=model_state,
-      mode=spec.ForwardPassMode.TRAIN,
-      rng=rng,
-      update_batch_norm=True)
+    params=current_model,
+    augmented_and_preprocessed_input_batch=batch,
+    model_state=model_state,
+    mode=spec.ForwardPassMode.TRAIN,
+    rng=rng,
+    update_batch_norm=True,
+  )
 
   label_smoothing = (
-      hyperparameters.label_smoothing if hasattr(HPARAMS,
-                                                 'label_smoothing') else 0.0)
+    hyperparameters.label_smoothing
+    if hasattr(HPARAMS, 'label_smoothing')
+    else 0.0
+  )
   if hasattr(hyperparameters, 'grad_clip'):
     grad_clip = hyperparameters.grad_clip
   else:
     grad_clip = None
 
   loss_dict = workload.loss_fn(
-      label_batch=batch['targets'],
-      logits_batch=logits_batch,
-      mask_batch=batch.get('weights'),
-      label_smoothing=label_smoothing)
+    label_batch=batch['targets'],
+    logits_batch=logits_batch,
+    mask_batch=batch.get('weights'),
+    label_smoothing=label_smoothing,
+  )
   summed_loss = loss_dict['summed']
   n_valid_examples = loss_dict['n_valid_examples']
   if USE_PYTORCH_DDP:
@@ -187,7 +197,8 @@ def update_params(
 
   if grad_clip is not None:
     torch.nn.utils.clip_grad_norm_(
-        current_model.parameters(), max_norm=grad_clip)
+      current_model.parameters(), max_norm=grad_clip
+    )
   optimizer_state['optimizer'].step()
   optimizer_state['scheduler'].step()
 
@@ -196,31 +207,38 @@ def update_params(
     with torch.no_grad():
       parameters = [p for p in current_model.parameters() if p.grad is not None]
       grad_norm = torch.norm(
-          torch.stack([torch.norm(p.grad.detach(), 2) for p in parameters]), 2)
+        torch.stack([torch.norm(p.grad.detach(), 2) for p in parameters]), 2
+      )
     if workload.metrics_logger is not None:
       workload.metrics_logger.append_scalar_metrics(
-          {
-              'loss': loss.item(),
-              'grad_norm': grad_norm.item(),
-          }, global_step)
-    logging.info('%d) loss = %0.3f, grad_norm = %0.3f',
-                 global_step,
-                 loss.item(),
-                 grad_norm.item())
+        {
+          'loss': loss.item(),
+          'grad_norm': grad_norm.item(),
+        },
+        global_step,
+      )
+    logging.info(
+      '%d) loss = %0.3f, grad_norm = %0.3f',
+      global_step,
+      loss.item(),
+      grad_norm.item(),
+    )
 
   return (optimizer_state, current_param_container, new_model_state)
 
 
-def prepare_for_eval(workload: spec.Workload,
-                     current_param_container: spec.ParameterContainer,
-                     current_params_types: spec.ParameterTypeTree,
-                     model_state: spec.ModelAuxiliaryState,
-                     hyperparameters: spec.Hyperparameters,
-                     loss_type: spec.LossType,
-                     optimizer_state: spec.OptimizerState,
-                     eval_results: List[Tuple[int, float]],
-                     global_step: int,
-                     rng: spec.RandomState) -> spec.UpdateReturn:
+def prepare_for_eval(
+  workload: spec.Workload,
+  current_param_container: spec.ParameterContainer,
+  current_params_types: spec.ParameterTypeTree,
+  model_state: spec.ModelAuxiliaryState,
+  hyperparameters: spec.Hyperparameters,
+  loss_type: spec.LossType,
+  optimizer_state: spec.OptimizerState,
+  eval_results: List[Tuple[int, float]],
+  global_step: int,
+  rng: spec.RandomState,
+) -> spec.UpdateReturn:
   """Return (updated_optimizer_state, updated_params)."""
   del workload
   del hyperparameters
@@ -234,8 +252,8 @@ def prepare_for_eval(workload: spec.Workload,
 
 def get_batch_size(workload_name):
   # Return the global batch size.
-  if hasattr(HPARAMS, "batch_size"):
-     return HPARAMS.batch_size
+  if hasattr(HPARAMS, 'batch_size'):
+    return HPARAMS.batch_size
   if workload_name == 'criteo1tb':
     return 262_144
   elif workload_name == 'fastmri':
@@ -262,14 +280,16 @@ def get_batch_size(workload_name):
     raise ValueError(f'Unsupported workload name: {workload_name}.')
 
 
-def data_selection(workload: spec.Workload,
-                   input_queue: Iterator[Dict[str, spec.Tensor]],
-                   optimizer_state: spec.OptimizerState,
-                   current_param_container: spec.ParameterContainer,
-                   model_state: spec.ModelAuxiliaryState,
-                   hyperparameters: spec.Hyperparameters,
-                   global_step: int,
-                   rng: spec.RandomState) -> Dict[str, spec.Tensor]:
+def data_selection(
+  workload: spec.Workload,
+  input_queue: Iterator[Dict[str, spec.Tensor]],
+  optimizer_state: spec.OptimizerState,
+  current_param_container: spec.ParameterContainer,
+  model_state: spec.ModelAuxiliaryState,
+  hyperparameters: spec.Hyperparameters,
+  global_step: int,
+  rng: spec.RandomState,
+) -> Dict[str, spec.Tensor]:
   """Select data from the infinitely repeating, pre-shuffled input queue.
   Each element of the queue is a batch of training examples and labels.
   """