Merge pull request #9 from borisdayma/feat--wandb-search

seq2seq/do_run.sh

@@ -6,5 +6,6 @@ python run_seq2seq_flax.py \
 	--per_device_train_batch_size 24 \
 	--per_device_eval_batch_size 24 \
 	--preprocessing_num_workers 48 \
+	--warmup_steps 1000 \
 	--do_train \
 	--do_eval \

seq2seq/run_seq2seq_flax.py

@@ -215,6 +215,13 @@ class DataTrainingArguments:
     overwrite_cache: bool = field(
         default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
     )
+    log_interval: Optional[int] = field(
+        default=5,
+        metadata={
+            "help": "For debugging purposes or quicker training, truncate the number of training examples to this "
+            "value if set."
+        },
+    )
 
     def __post_init__(self):
         if self.dataset_name is None and self.train_file is None and self.validation_file is None:
@@ -307,12 +314,12 @@ def write_metric(summary_writer, train_metrics, eval_metrics, train_time, step):
 
     train_metrics = get_metrics(train_metrics)
     for key, vals in train_metrics.items():
-        tag = f"train_{key}"
+        tag = f"train_epoch/{key}"
         for i, val in enumerate(vals):
             summary_writer.scalar(tag, val, step - len(vals) + i + 1)
 
     for metric_name, value in eval_metrics.items():
-        summary_writer.scalar(f"eval_{metric_name}", value, step)
+        summary_writer.scalar(f"eval/{metric_name}", value, step)
 
 
 def create_learning_rate_fn(
@@ -616,17 +623,24 @@ def main():
         return traverse_util.unflatten_dict(flat_mask)
 
     # create adam optimizer
-    adamw = optax.adamw(
-        learning_rate=linear_decay_lr_schedule_fn,
-        b1=training_args.adam_beta1,
-        b2=training_args.adam_beta2,
-        eps=training_args.adam_epsilon,
-        weight_decay=training_args.weight_decay,
-        mask=decay_mask_fn,
-    )
+    if training_args.adafactor:
+        # We use the default parameters here to initialize adafactor,
+        # For more details about the parameters please check https://github.com/deepmind/optax/blob/ed02befef9bf81cbbf236be3d2b0e032e9ed4a40/optax/_src/alias.py#L74
+        optimizer = optax.adafactor(
+            learning_rate=linear_decay_lr_schedule_fn,
+        )
+    else:
+        optimizer = optax.adamw(
+            learning_rate=linear_decay_lr_schedule_fn,
+            b1=training_args.adam_beta1,
+            b2=training_args.adam_beta2,
+            eps=training_args.adam_epsilon,
+            weight_decay=training_args.weight_decay,
+            mask=decay_mask_fn,
+        )
 
     # Setup train state
-    state = TrainState.create(apply_fn=model.__call__, params=model.params, tx=adamw, dropout_rng=dropout_rng)
+    state = TrainState.create(apply_fn=model.__call__, params=model.params, tx=optimizer, dropout_rng=dropout_rng)
 
     # label smoothed cross entropy
     def loss_fn(logits, labels, padding_mask, label_smoothing_factor=0.0):
@@ -718,6 +732,7 @@ def main():
 
     train_time = 0
     epochs = tqdm(range(num_epochs), desc=f"Epoch ... (1/{num_epochs})", position=0)
+    global_step = 0
     for epoch in epochs:
         # ======================== Training ================================
         train_start = time.time()
@@ -730,11 +745,16 @@ def main():
         train_loader = data_loader(input_rng, train_dataset, train_batch_size, shuffle=True)
         steps_per_epoch = len(train_dataset) // train_batch_size
         # train
-        for _ in tqdm(range(steps_per_epoch), desc="Training...", position=1, leave=False):
+        for step in tqdm(range(steps_per_epoch), desc="Training...", position=1, leave=False):
+            global_step +=1
             batch = next(train_loader)
             state, train_metric = p_train_step(state, batch)
             train_metrics.append(train_metric)
 
+            if global_step % data_args.log_interval == 0 and jax.process_index() == 0:
+                for k, v in unreplicate(train_metric).items():
+                    wandb.log(f{'train/{k}': jax.device_get(v)}, step=global_step)
+
         train_time += time.time() - train_start
 
         train_metric = unreplicate(train_metric)