Plots results (#37)

Author: nathanielsimard

environment.yml

@@ -11,6 +11,7 @@ dependencies:
   - cudatoolkit =10.1
   - plotly >=4.8
   - plotly-orca >=1.3
+  - matplotlib
   - numpy
   - pandas
   - pyyaml

mcp/argparser.py

@@ -32,6 +32,10 @@ def parse_eval_arguments():
     return parser.parse_args()
 
 
+# For now both are the same
+parse_viz_arguments = parse_eval_arguments
+
+
 def _default_parser():
     parser = argparse.ArgumentParser()
     parser.add_argument(

mcp/context/base.py

@@ -40,6 +40,7 @@
 from mcp.task.supervised import SupervisedTask
 from mcp.training.loop import TrainingLoop
 from mcp.training.trainer import Trainer, TrainerLoggers
+from mcp.viz.base import Vizualization
 
 ValidFewShotDataLoaderFactory = NewType(
     "ValidFewShotDataLoaderFactory", FewShotDataLoaderFactory
@@ -96,10 +97,10 @@ def provide_evaluation_loggers(self) -> EvaluationLoggers:
 
         return EvaluationLoggers(
             support=ResultLogger(
-                "Evaluation Support", os.path.join(output_dir, "support")
+                "Evaluation Support", os.path.join(output_dir, "test-support")
             ),
             evaluation=ResultLogger(
-                "Evaluation Support", os.path.join(output_dir, "eval")
+                "Evaluation Support", os.path.join(output_dir, "test-eval")
             ),
         )
 
@@ -136,6 +137,12 @@ def provide_evaluation(
     def provide_experiment_result(self) -> ExperimentResult:
         return ExperimentResult(self.config, self.output_dir)
 
+    @provider
+    @inject
+    @singleton
+    def provide_vizualization(self, result: ExperimentResult) -> Vizualization:
+        return Vizualization(result)
+
 
 class TrainerModule(Module):
     def __init__(self, config: ExperimentConfig, output_dir: str, device: torch.device):

mcp/main.py

@@ -7,6 +7,7 @@
 from mcp.evaluation import Evaluation
 from mcp.result.experiment import ExperimentResult
 from mcp.training.trainer import Trainer
+from mcp.viz.base import Vizualization
 
 
 def run_train(
@@ -33,3 +34,11 @@ def run_eval(config: ExperimentConfig, result_dir: str, device_str: str):
     evaluation = injector.get(Evaluation)
     result = injector.get(ExperimentResult)
     evaluation.eval(result.best_epoch())
+
+
+def run_viz(config: ExperimentConfig, result_dir: str, device_str: str):
+    device = torch.device(device_str)
+    injector = create_injector(config, result_dir, device)
+
+    viz = injector.get(Vizualization)
+    viz.plot()

mcp/result/experiment.py

@@ -1,5 +1,6 @@
 import os
-from typing import List
+import sys
+from typing import Callable, List, Optional
 
 import numpy as np
 
@@ -10,27 +11,56 @@
 logger = create_logger(__name__)
 
 
+class EpochResult(object):
+    def __init__(self, file_name: str):
+        self.file_name = file_name
+
+    def load(self) -> List[List[ResultRecord]]:
+        return load_records_from_file(self.file_name)
+
+    @staticmethod
+    def losses(records: List[List[ResultRecord]]) -> List[List[float]]:
+        return [[r.loss for r in rec] for rec in records]
+
+    @staticmethod
+    def metric(records: List[List[ResultRecord]]) -> List[List[float]]:
+        return [[r.metric for r in rec] for rec in records]
+
+    @staticmethod
+    def task_name(records: List[List[ResultRecord]]) -> List[str]:
+        return [r.name for r in records[0]]
+
+    @staticmethod
+    def metric_name(records: List[List[ResultRecord]]) -> List[str]:
+        return [r.metric_name for r in records[0]]
+
+    @staticmethod
+    def reduce(
+        values: List[List[float]],
+        reduce_task: Optional[Callable] = np.mean,
+        reduce_iter: Optional[Callable] = np.mean,
+    ) -> np.ndarray:
+        if reduce_task is None and reduce_iter is None:
+            raise ValueError("Must reduce on something")
+
+        if reduce_task is not None:
+            values = [reduce_task(np.asarray(vv), axis=-1) for vv in values]
+
+        if reduce_iter is not None:
+            values = reduce_iter(np.asarray(values), axis=0)
+
+        return values
+
+
 class ExperimentResult(object):
     def __init__(self, config: ExperimentConfig, output_dir: str):
         self.config = config
         self.output_dir = output_dir
-        self._records_dir = os.path.join(self.output_dir, "train")
+        self._records_dir_train = os.path.join(self.output_dir, "train")
+        self._records_dir_eval = os.path.join(self.output_dir, "evaluation")
 
     def best_epoch(self) -> int:
-        losses = []
-        for epoch in range(1, self.config.trainer.epochs + 1):
-            try:
-                file_name = os.path.join(self._records_dir, f"eval-{epoch}")
-                records_valid = load_records_from_file(file_name)
-                loss = np.asarray(
-                    [self._records_loss(rs) for rs in records_valid]
-                ).mean()
-                losses.append(loss)
-            except FileNotFoundError:
-                logger.warning(
-                    f"Training did not complete {epoch-1}/{self.config.trainer.epochs}"
-                )
-                break
+        losses = self.metric("train", EpochResult.losses)
 
         indexes = np.argsort(np.asarray(losses))
         index = indexes[0]
@@ -40,5 +70,38 @@ def best_epoch(self) -> int:
         logger.info(f"Found the best epoch to be {epoch} with valid loss {valid_loss}")
         return epoch
 
-    def _records_loss(self, records: List[ResultRecord]) -> float:
-        return np.asarray([r.loss for r in records]).mean()
+    def records(self, tag: str, train: bool = True) -> List[EpochResult]:
+        records_dir = self._records_dir_train if train else self._records_dir_eval
+
+        results: List[EpochResult] = []
+        for epoch in range(1, sys.maxsize):
+            file_name = os.path.join(records_dir, f"{tag}-{epoch}")
+            if not os.path.exists(file_name):
+                break
+
+            results.append(EpochResult(file_name))
+
+        return results
+
+    def task_names(self, tag: str, train: bool = True) -> List[str]:
+        e_records = self.records(tag, train=train)[0]
+        return EpochResult.task_name(e_records.load())
+
+    def metric_names(self, tag: str, train: bool = True) -> List[str]:
+        e_records = self.records(tag, train)[0]
+        return EpochResult.metric_name(e_records.load())
+
+    def metric(
+        self, tag: str, metric, reduce_task=np.mean, reduce_iter=np.mean, train=True
+    ) -> np.ndarray:
+        e_records = self.records(tag, train=train)
+        return np.asarray(
+            [
+                EpochResult.reduce(
+                    metric(records.load()),
+                    reduce_task=reduce_task,
+                    reduce_iter=reduce_iter,
+                )
+                for records in e_records
+            ]
+        )

mcp/result/logger.py

@@ -33,6 +33,10 @@ def load_records_from_file(file_path: str) -> List[List[ResultRecord]]:
 
 
 def load_records(line: str) -> List[ResultRecord]:
+    """Load all records for an iteration.
+
+    The number of records is determined by the number of tasks.
+    """
     objs = json.loads(line)
     return [
         ResultRecord(

mcp/viz/__init__.py

@@ -0,0 +1,17 @@
+import os
+
+from mcp.result.experiment import ExperimentResult
+from mcp.viz.loss import plot_loss
+from mcp.viz.metric import plot_metric
+
+
+class Vizualization(object):
+    def __init__(self, results: ExperimentResult):
+        self.results = results
+
+    def plot(self):
+        output_dir = os.path.join(self.results.output_dir, "viz")
+        os.makedirs(output_dir, exist_ok=True)
+
+        plot_loss(output_dir, self.results)
+        plot_metric(output_dir, self.results)

mcp/viz/base.py

@@ -0,0 +1,43 @@
+from typing import List
+
+import numpy as np
+from matplotlib import pyplot as plt
+from matplotlib.ticker import MaxNLocator
+
+
+def line_plot(
+    task_names_train: List[str],
+    task_names_eval: List[str],
+    values_train: np.ndarray,
+    values_eval: np.ndarray,
+    y_label: str,
+    x_label: str = "Epoch",
+    bbox_to_anchor=(0.90, 0.88),
+    y_int: bool = False,
+    x_int: bool = True,
+) -> plt.Figure:
+
+    fig = plt.figure()
+    ax = fig.subplots()
+    for i, name in enumerate(task_names_train):
+        x = list(range(len(values_train)))
+        ax.plot(
+            x, values_train[:, i], label=f"Train - {name}",
+        )
+
+    for i, name in enumerate(task_names_eval):
+        x = list(range(len(values_eval)))
+        ax.plot(
+            x, values_eval[:, i], label=f"Valid - {name}", linestyle=":",
+        )
+
+    ax.set_xlabel(x_label)
+    ax.set_ylabel(y_label)
+
+    if x_int:
+        ax.xaxis.set_major_locator(MaxNLocator(integer=True))
+    if y_int:
+        ax.yaxis.set_major_locator(MaxNLocator(integer=True))
+
+    fig.legend(bbox_to_anchor=bbox_to_anchor)
+    return fig

mcp/viz/line_plot.py

@@ -0,0 +1,19 @@
+import os
+
+from mcp.result.experiment import EpochResult, ExperimentResult
+from mcp.viz.line_plot import line_plot
+
+
+def plot_loss(output_dir: str, results: ExperimentResult):
+    losses_train = results.metric("train", EpochResult.losses, reduce_task=None)
+    task_names_train = results.task_names("train")
+
+    losses_eval = results.metric("eval", EpochResult.losses, reduce_task=None)
+    task_names_eval = results.task_names("eval")
+
+    fig = line_plot(
+        task_names_train, task_names_eval, losses_train, losses_eval, "Loss"
+    )
+
+    file_name = os.path.join(output_dir, "losses.png")
+    fig.savefig(file_name)

mcp/viz/loss.py

@@ -0,0 +1,44 @@
+import os
+
+import numpy as np
+
+from mcp.result.experiment import EpochResult, ExperimentResult
+from mcp.utils import logging
+from mcp.viz.line_plot import line_plot
+
+logger = logging.create_logger(__name__)
+
+
+def plot_metric(output_dir: str, results: ExperimentResult):
+    metric_train = results.metric("train", EpochResult.metric, reduce_task=None)
+    task_names_train = results.task_names("train")
+
+    metric_eval = results.metric("eval", EpochResult.metric, reduce_task=None)
+    task_names_eval = results.task_names("eval")
+
+    metric_names_train = results.metric_names("train")
+    metric_names_eval = results.metric_names("eval")
+
+    for i, (task, metric) in enumerate(zip(task_names_train, metric_names_train)):
+        fig = line_plot([task], [], metric_train[:, i : i + 1], np.array([]), metric)
+        file_name = os.path.join(output_dir, f"metric-{task}-{metric}-train.png")
+        fig.savefig(file_name)
+
+    for i, (task, metric) in enumerate(zip(task_names_eval, metric_names_eval)):
+        fig = line_plot([], [task], np.array([]), metric_eval[:, i : i + 1], metric)
+        file_name = os.path.join(output_dir, f"metric-{task}-{metric}-eval.png")
+        fig.savefig(file_name)
+
+    metric_test_eval = results.metric("eval", EpochResult.metric, train=False)
+
+    if len(metric_test_eval) > 0:
+        metric_name_test_eval = results.metric_names("eval", train=False)[0]
+        task_name_test_eval = results.task_names("eval", train=False)[0]
+
+        values = np.asarray(metric_test_eval)
+        mean = np.mean(values)
+        std = np.std(values, ddof=1)
+
+        logger.info(
+            f"Test {task_name_test_eval}: {mean} +- {std} {metric_name_test_eval}"
+        )

mcp/viz/metric.py

@@ -0,0 +1,22 @@
+#!/usr/bin/env python
+import os
+
+from mcp.argparser import initialize_logging, parse_viz_arguments
+
+
+def run(args):
+    from mcp import main
+    from mcp.config.loader import load
+    from mcp.config.parser import parse
+
+    config_path = os.path.join(args.result, "config_full.yml")
+    configs = load(config_path)
+    config_experiment = parse([configs])
+
+    main.run_viz(config_experiment, args.result, args.device)  # type: ignore
+
+
+if __name__ == "__main__":
+    args = parse_viz_arguments()
+    initialize_logging(args.logging, args.result, args.debug)
+    run(args)