Refactored evaluation fun and renamed variable names to generic names.

Author: quic-meetkuma

QEfficient/finetune/utils/train_utils.py

@@ -62,9 +62,9 @@ def train(
 
     Returns: results dictionary containing average training and validation perplexity and loss
     """
-    train_prep = []
+    train_metric = []
     train_loss = []
-    val_prep = []
+    val_metric = []
     val_loss = []
 
     if train_config.save_metrics:
@@ -73,10 +73,10 @@ def train(
         metrics_filename = (
             f"{train_config.output_dir}/metrics_data_{local_rank}-{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}.json"
         )
-        train_step_perplexity = []
+        train_step_metric = []
         train_step_loss = []
         val_step_loss = []
-        val_step_perplexity = []
+        val_step_metric = []
 
     epoch_times = []
     checkpoint_times = []
@@ -106,10 +106,10 @@ def train(
 
     acc_helper = None
     if train_config.task_type == "seq_classification":
-        if local_rank is None:
-            num_classes = model.classifier.out_features
-        else:
+        if train_config.enable_ddp:
             num_classes = model.module.classifier.out_features
+        else:
+            num_classes = model.classifier.out_features
         acc_helper = torchmetrics.classification.MulticlassAccuracy(num_classes=num_classes).to(device)
 
     # Start the training loop
@@ -231,7 +231,11 @@ def train(
 
             if train_config.save_metrics:
                 train_step_loss.append(loss.detach().float().item())
-                train_step_perplexity.append(float(torch.exp(loss.detach().float())))
+                if train_config.task_type == "seq_classification":
+                    step_metric_val = acc_helper.compute()
+                else:
+                    step_metric_val = float(torch.exp(loss.detach().float()))
+                train_step_metric.append(step_metric_val)
 
             if train_config.grad_scaler:
                 scaler.scale(loss).backward()  # backward pass
@@ -266,12 +270,12 @@ def train(
                     metrics_filename,
                     train_step_loss,
                     train_loss,
-                    train_step_perplexity,
-                    train_prep,
+                    train_step_metric,
+                    train_metric,
                     val_step_loss,
                     val_loss,
-                    val_step_perplexity,
-                    val_prep,
+                    val_step_metric,
+                    val_metric,
                 )
             if train_config.enable_ddp:
                 if loss_0_counter.item() == train_config.convergence_counter:
@@ -307,11 +311,11 @@ def train(
             if train_config.enable_ddp:
                 dist.all_reduce(accuracy, op=dist.ReduceOp.SUM)
                 accuracy /= dist.get_world_size()
-            train_perplexity = accuracy
+            train_metric = accuracy
         else:
-            train_perplexity = torch.exp(train_epoch_loss)
+            train_metric = torch.exp(train_epoch_loss)
 
-        train_prep.append(float(train_perplexity))
+        train_metric.append(float(train_metric))
         train_loss.append(float(train_epoch_loss))
 
         # Update the learning rate as needed
@@ -320,21 +324,21 @@ def train(
         if train_config.run_validation:
             if train_config.enable_ddp:
                 dist.barrier()
-                eval_ppl, eval_epoch_loss, temp_val_loss, temp_step_perplexity = evaluation_helper(
-                    model, train_config, eval_dataloader, local_rank, tokenizer, device
+                eval_epoch_loss, eval_metric, temp_val_loss, temp_step_metric = evaluation_helper(
+                    model, train_config, eval_dataloader, device
                 )
                 if local_rank == 0:
                     tensorboard_updates.add_scalars("loss", {"eval": eval_epoch_loss}, total_train_steps)
 
             else:
-                eval_ppl, eval_epoch_loss, temp_val_loss, temp_step_perplexity = evaluation_helper(
-                    model, train_config, eval_dataloader, local_rank, tokenizer, device
+                eval_epoch_loss, eval_metric, temp_val_loss, temp_step_metric = evaluation_helper(
+                    model, train_config, eval_dataloader, device
                 )
                 tensorboard_updates.add_scalars("loss", {"eval": eval_epoch_loss}, total_train_steps)
 
             if train_config.save_metrics:
                 val_step_loss.extend(temp_val_loss)
-                val_step_perplexity.extend(temp_step_perplexity)
+                val_step_metric.extend(temp_step_metric)
 
         # saving the adapters after completion of each epoch
         if train_config.save_model:
@@ -349,14 +353,14 @@ def train(
                 best_val_loss = eval_epoch_loss
                 print(f"best eval loss on epoch {epoch + 1} is {best_val_loss}")
             val_loss.append(float(eval_epoch_loss))
-            val_prep.append(float(eval_ppl))
+            val_metric.append(float(eval_metric))
         if train_config.task_type == "seq_classification":
             print(
-                f"Epoch {epoch + 1}: train_acc={train_perplexity:.4f}, train_epoch_loss={train_epoch_loss:.4f}, epoch time {epoch_end_time}s"
+                f"Epoch {epoch + 1}: train_acc={train_metric:.4f}, train_epoch_loss={train_epoch_loss:.4f}, epoch time {epoch_end_time}s"
             )
         else:
             print(
-                f"Epoch {epoch + 1}: train_perplexity={train_perplexity:.4f}, train_epoch_loss={train_epoch_loss:.4f}, epoch time {epoch_end_time}s"
+                f"Epoch {epoch + 1}: train_metric={train_metric:.4f}, train_epoch_loss={train_epoch_loss:.4f}, epoch time {epoch_end_time}s"
             )
 
         # Saving the results every epoch to plot later
@@ -365,31 +369,25 @@ def train(
                 metrics_filename,
                 train_step_loss,
                 train_loss,
-                train_step_perplexity,
-                train_prep,
+                train_step_metric,
+                train_metric,
                 val_step_loss,
                 val_loss,
-                val_step_perplexity,
-                val_prep,
+                val_step_metric,
+                val_metric,
             )
     avg_epoch_time = sum(epoch_times) / len(epoch_times)
     avg_checkpoint_time = sum(checkpoint_times) / len(checkpoint_times) if len(checkpoint_times) > 0 else 0
-    avg_train_prep = sum(train_prep) / len(train_prep)
+    avg_train_metric = sum(train_metric) / len(train_metric)
     avg_train_loss = sum(train_loss) / len(train_loss)
     if train_config.run_validation:
-        avg_eval_prep = sum(val_prep) / len(val_prep)
+        avg_eval_metric = sum(val_metric) / len(val_metric)
         avg_eval_loss = sum(val_loss) / len(val_loss)
 
-    if train_config.task_type == "seq_classification":
-        results["avg_train_acc"] = avg_train_prep
-    else:
-        results["avg_train_prep"] = avg_train_prep
+    results["avg_train_metric"] = avg_train_metric
     results["avg_train_loss"] = avg_train_loss
     if train_config.run_validation:
-        if train_config.task_type == "seq_classification":
-            results["avg_eval_acc"] = avg_eval_prep
-        else:
-            results["avg_eval_prep"] = avg_eval_prep
+        results["avg_eval_metric"] = avg_eval_metric
         results["avg_eval_loss"] = avg_eval_loss
     results["avg_epoch_time"] = avg_epoch_time
     results["avg_checkpoint_time"] = avg_checkpoint_time
@@ -399,39 +397,40 @@ def train(
     return results
 
 
-def evaluation_ppl(model, train_config, eval_dataloader, local_rank, tokenizer, device):
+def evaluation_helper(model, train_config, eval_dataloader, device):
     """
     Evaluates the model on the given dataloader
 
     Args:
         model: The model to evaluate
         eval_dataloader: The dataloader containing the evaluation data
-        local_rank: The rank of the current node in a distributed setting
-        tokenizer: The tokenizer used to decode predictions
 
-    Returns: eval_ppl, eval_epoch_loss
+    Returns: eval_epoch_loss, eval_metric, eval_step_loss, eval_step_metric
     """
     model.eval()
 
+    if train_config.task_type == "seq_classification":
+        if train_config.enable_ddp:
+            num_classes = model.module.classifier.out_features
+        else:
+            num_classes = model.classifier.out_features
+        acc_helper = torchmetrics.classification.MulticlassAccuracy(num_classes=num_classes).to(device)
+
     # special handling for qaic device and dtype
     # model.to(device)
 
-    eval_preds = []
     val_step_loss = []
-    val_step_perplexity = []
+    val_step_metric = []
 
     eval_loss = 0.0  # Initialize evaluation loss
-    total_eval_steps = 0
-    # max_steps_reached = False  # Flag to indicate max eval steps reached
 
     for step, batch in enumerate(tqdm(eval_dataloader, colour="green", desc="evaluating Epoch", dynamic_ncols=True)):
-        total_eval_steps += 1
         #  stop when the maximum number of eval steps is reached
-        if train_config.max_eval_step > 0 and total_eval_steps > train_config.max_eval_step:
-            # max_steps_reached = True
+        if train_config.max_eval_step > 0 and step > train_config.max_eval_step:
             break
         for key in batch.keys():
             batch[key] = batch[key].to(device)
+
         # Ensure no gradients are computed for this scope to save memory
         with torch.no_grad():
             # Forward pass and compute loss
@@ -441,100 +440,32 @@ def evaluation_ppl(model, train_config, eval_dataloader, local_rank, tokenizer,
                 outputs = model(**batch)
             loss = outputs.loss
 
-            if train_config.save_metrics:
-                val_step_loss.append(loss.detach().float().item())
-                val_step_perplexity.append(float(torch.exp(loss.detach().float())))
-
-            eval_loss += loss.detach().float()
-        # Decode predictions and add to evaluation predictions list
-        preds = torch.argmax(outputs.logits, -1)
-        eval_preds.extend(tokenizer.batch_decode(preds.detach().cpu().numpy(), skip_special_tokens=True))
-
-    # Compute average loss and perplexity
-    eval_epoch_loss = eval_loss / len(eval_dataloader)
-    eval_ppl = torch.exp(eval_epoch_loss)
-
-    # Print evaluation metrics
-    print(f" {eval_ppl.detach().cpu()=} {eval_epoch_loss.detach().cpu()=}")
-
-    return eval_ppl, eval_epoch_loss, val_step_loss, val_step_perplexity
-
-
-def evaluation_acc(model, train_config, eval_dataloader, local_rank, tokenizer, device):
-    """
-    Evaluates the model on the given dataloader
-
-    Args:
-        model: The model to evaluate
-        eval_dataloader: The dataloader containing the evaluation data
-        local_rank: The rank of the current node in a distributed setting
-        tokenizer: The tokenizer used to decode predictions
-
-    Returns: eval_acc, eval_epoch_loss
-    """
-    model.eval()
-    if train_config.enable_ddp:
-        num_classes = model.module.classifier.out_features
-    else:
-        num_classes = model.classifier.out_features
-
-    acc_helper = torchmetrics.classification.MulticlassAccuracy(num_classes=num_classes).to(device)
-
-    # special handling for qaic device and dtype
-    # model.to(device)
-
-    # eval_preds = []
-    val_step_loss = []
-    val_step_acc = []
-
-    eval_loss = 0.0  # Initialize evaluation loss
-    total_eval_steps = 0
-    # max_steps_reached = False  # Flag to indicate max eval steps reached
+            if train_config.task_type == "seq_classification":
+                logits = outputs.logits
+                labels = batch["labels"][:, 0]
+                preds = torch.nn.functional.softmax(logits, dim=-1)
+                val_acc = acc_helper.forward(preds, labels)
+                metric_val = val_acc.detach().float().item()
+            else:
+                metric_val = float(torch.exp(loss.detach().float()))
 
-    for step, batch in enumerate(tqdm(eval_dataloader, colour="green", desc="evaluating Epoch", dynamic_ncols=True)):
-        total_eval_steps += 1
-        #  stop when the maximum number of eval steps is reached
-        if train_config.max_eval_step > 0 and total_eval_steps > train_config.max_eval_step:
-            # max_steps_reached = True
-            break
-        for key in batch.keys():
-            batch[key] = batch[key].to(device)
-        # Ensure no gradients are computed for this scope to save memory
-        with torch.no_grad():
-            # Forward pass and compute loss
-            with (
-                torch.autocast(device_type=device, dtype=torch.float16) if train_config.use_autocast else nullcontext()
-            ):
-                outputs = model(**batch)
-            loss = outputs.loss
-            logits = outputs.logits
-            labels = batch["labels"][:, 0]
             if train_config.save_metrics:
                 val_step_loss.append(loss.detach().float().item())
-                preds = torch.nn.functional.softmax(logits, dim=-1)
-                val_acc = acc_helper.forward(preds, labels)
-                val_step_acc.append(val_acc.detach().float().item())
+                val_step_metric.append(metric_val)
 
             eval_loss += loss.detach().float()
-        # Decode predictions and add to evaluation predictions list
-        # preds = torch.argmax(outputs.logits, -1)
-        # eval_preds.extend(tokenizer.batch_decode(preds.detach().cpu().numpy(), skip_special_tokens=True))
 
-    # Compute average loss and perplexity
+    # Compute average loss and metric
     eval_epoch_loss = eval_loss / len(eval_dataloader)
-    eval_acc = acc_helper.compute()
+    if train_config.task_type == "seq_classification":
+        eval_metric = acc_helper.compute()
+    else:
+        eval_metric = torch.exp(eval_epoch_loss)
 
     # Print evaluation metrics
-    print(f" {eval_acc.detach().cpu()=} {eval_epoch_loss.detach().cpu()=}")
+    print(f" {eval_metric.detach().cpu()=} {eval_epoch_loss.detach().cpu()=}")
 
-    return eval_acc, eval_epoch_loss, val_step_loss, val_step_acc
-
-
-def evaluation_helper(model, train_config, eval_dataloader, local_rank, tokenizer, device):
-    if train_config.task_type == "seq_classification":
-        return evaluation_acc(model, train_config, eval_dataloader, local_rank, tokenizer, device)
-    else:
-        return evaluation_ppl(model, train_config, eval_dataloader, local_rank, tokenizer, device)
+    return eval_metric, eval_epoch_loss, val_step_loss, val_step_metric
 
 
 def get_longest_seq_length(data: List[Dict]) -> Tuple[int, int]:
@@ -571,22 +502,22 @@ def save_to_json(
     output_filename,
     train_step_loss,
     train_epoch_loss,
-    train_step_ppl,
-    train_epoch_ppl,
+    train_step_metric,
+    train_epoch_metric,
     val_step_loss,
     val_epoch_loss,
-    val_step_ppl,
-    val_epoch_ppl,
+    val_step_metric,
+    val_epoch_metric,
 ):
     metrics_data = {
         "train_step_loss": train_step_loss,
         "train_epoch_loss": train_epoch_loss,
-        "train_step_perplexity": train_step_ppl,
-        "train_epoch_perplexity": train_epoch_ppl,
+        "train_step_metric": train_step_metric,
+        "train_epoch_metric": train_epoch_metric,
         "val_step_loss": val_step_loss,
         "val_epoch_loss": val_epoch_loss,
-        "val_step_perplexity": val_step_ppl,
-        "val_epoch_perplexity": val_epoch_ppl,
+        "val_step_metric": val_step_metric,
+        "val_epoch_metric": val_epoch_metric,
     }
     with open(output_filename, "w") as f:
         json.dump(metrics_data, f)