fix multiple cuda error, save evals to right folder + cleanup

elk/evaluation/evaluate.py

@@ -2,21 +2,17 @@
 from functools import partial
 from pathlib import Path
 from typing import Callable, Literal, Optional
-from dataclasses import dataclass
-from functools import partial
-from pathlib import Path
-from typing import Callable, Literal, Optional
 
+import pandas as pd
 import torch
 from simple_parsing.helpers import Serializable, field
 
-from evaluation.evaluate_log import EvalLog
-from extraction.extraction import Extract
-from files import elk_reporter_dir
-from run import Run
-from training import Reporter
-from training.baseline import evaluate_baseline, load_baseline, train_baseline
-from utils import select_usable_devices
+from ..extraction.extraction import Extract
+from ..files import create_output_directory, elk_reporter_dir, memorably_named_dir
+from ..run import Run
+from ..training import Reporter
+from ..training.baseline import evaluate_baseline, load_baseline
+from ..utils import select_usable_devices
 
 
 @dataclass
@@ -48,8 +44,9 @@ class Eval(Serializable):
 
  def execute(self):
  transfer_eval = elk_reporter_dir() / self.source / "transfer_eval"
+ out_dir = memorably_named_dir(transfer_eval)
 
- run = Evaluate(cfg=self, out_dir=transfer_eval)
+ run = Evaluate(cfg=self, out_dir=out_dir)
  run.evaluate()
 
 
@@ -59,7 +56,7 @@ class Evaluate(Run):
 
  def evaluate_reporter(
  self, layer: int, devices: list[str], world_size: int = 1
- ) -> EvalLog:
+ ) -> pd.Series:
  """Evaluate a single reporter on a single layer."""
  device = self.get_device(devices, world_size)
 
@@ -70,9 +67,7 @@ def evaluate_reporter(
 
  experiment_dir = elk_reporter_dir() / self.cfg.source
 
- reporter_path = (
- experiment_dir / "reporters" / f"layer_{layer}.pt"
- )
+ reporter_path = experiment_dir / "reporters" / f"layer_{layer}.pt"
  reporter: Reporter = torch.load(reporter_path, map_location=device)
  reporter.eval()
 
@@ -82,129 +77,25 @@ def evaluate_reporter(
  test_x1,
  )
 
- lr_dir = experiment_dir / "lr_models"
- if not self.cfg.skip_baseline and lr_dir.exists():
- lr_model = load_baseline(lr_dir, layer)
- lr_auroc, lr_acc = evaluate_baseline(lr_model, test_x0, test_x1, test_labels)
-
- print("transfer_eval", lr_auroc, lr_acc)
-
- # stats.lr_auroc = lr_auroc
- # stats.lr_acc = lr_acc
- # save_baseline(lr_dir, layer, lr_model)
-
- return EvalLog(
- layer=layer,
- eval_result=test_result,
- )
-
- def evaluate(self):
- """Evaluate the reporter on all layers."""
- devices = select_usable_devices(
- self.cfg.num_gpus, min_memory=self.cfg.data.min_gpu_mem
- )
-
- num_devices = len(devices)
- func: Callable[[int], EvalLog] = partial(
- self.evaluate_reporter, devices=devices, world_size=num_devices
- )
- self.apply_to_layers(
- func=func,
- num_devices=num_devices,
- to_csv_line=lambda item: item.to_csv_line(),
- csv_columns=EvalLog.csv_columns(),
- )
-
-import torch
-from simple_parsing.helpers import Serializable, field
-
-from evaluation.evaluate_log import EvalLog
-from extraction.extraction import Extract
-from files import elk_reporter_dir
-from run import Run
-from training import Reporter
-from training.baseline import evaluate_baseline, load_baseline, train_baseline
-from utils import select_usable_devices
-
-
-@dataclass
-class Eval(Serializable):
- """
- Full specification of a reporter evaluation run.
-
- Args:
- data: Config specifying hidden states on which the reporter will be evaluated.
- source: The name of the source run directory
- which contains the reporters directory.
- normalization: The normalization method to use. Defaults to "meanonly". See
- `elk.training.preprocessing.normalize()` for details.
- num_gpus: The number of GPUs to use. Defaults to -1, which means
- "use all available GPUs".
- debug: When in debug mode, a useful log file is saved to the memorably-named
- output directory. Defaults to False.
- """
-
- data: Extract
- source: str = field(positional=True)
- normalization: Literal["legacy", "none", "elementwise", "meanonly"] = "meanonly"
-
- debug: bool = False
- out_dir: Optional[Path] = None
- num_gpus: int = -1
- skip_baseline: bool = False
- concatenated_layer_offset: int = 0
-
- def execute(self):
- transfer_eval = elk_reporter_dir() / self.source / "transfer_eval"
-
- run = Evaluate(cfg=self, out_dir=transfer_eval)
- run.evaluate()
-
-
-@dataclass
-class Evaluate(Run):
- cfg: Eval
-
- def evaluate_reporter(
- self, layer: int, devices: list[str], world_size: int = 1
- ) -> EvalLog:
- """Evaluate a single reporter on a single layer."""
- device = self.get_device(devices, world_size)
-
- _, _, test_x0, test_x1, _, test_labels = self.prepare_data(
- device,
- layer,
- )
-
- experiment_dir = elk_reporter_dir() / self.cfg.source
-
- reporter_path = (
- experiment_dir / "reporters" / f"layer_{layer}.pt"
- )
- reporter: Reporter = torch.load(reporter_path, map_location=device)
- reporter.eval()
-
- test_result = reporter.score(
- test_labels,
- test_x0,
- test_x1,
+ stats_row = pd.Series(
+ {
+ "layer": layer,
+ **test_result._asdict(),
+ }
  )
 
  lr_dir = experiment_dir / "lr_models"
  if not self.cfg.skip_baseline and lr_dir.exists():
  lr_model = load_baseline(lr_dir, layer)
- lr_auroc, lr_acc = evaluate_baseline(lr_model, test_x0, test_x1, test_labels)
+ lr_model.eval()
+ lr_auroc, lr_acc = evaluate_baseline(
+ lr_model.cuda(), test_x0.cuda(), test_x1.cuda(), test_labels
+ )
 
- print("transfer_eval", lr_auroc, lr_acc)
+ stats_row["lr_auroc"] = lr_auroc
+ stats_row["lr_acc"] = lr_acc
 
- # stats.lr_auroc = lr_auroc
- # stats.lr_acc = lr_acc
- # save_baseline(lr_dir, layer, lr_model)
-
- return EvalLog(
- layer=layer,
- eval_result=test_result,
- )
+ return stats_row
 
  def evaluate(self):
  """Evaluate the reporter on all layers."""
@@ -213,12 +104,7 @@ def evaluate(self):
  )
 
  num_devices = len(devices)
- func: Callable[[int], EvalLog] = partial(
+ func: Callable[[int], pd.Series] = partial(
  self.evaluate_reporter, devices=devices, world_size=num_devices
  )
- self.apply_to_layers(
- func=func,
- num_devices=num_devices,
- to_csv_line=lambda item: item.to_csv_line(),
- csv_columns=EvalLog.csv_columns(),
- )
+ self.apply_to_layers(func=func, num_devices=num_devices)

elk/run.py

@@ -6,16 +6,17 @@
 from typing import TYPE_CHECKING, Callable, Iterator, Optional, Union
 
 import numpy as np
+import pandas as pd
 import torch
 import torch.multiprocessing as mp
+from datasets import DatasetDict
 from torch import Tensor
 from tqdm import tqdm
 
-from datasets import DatasetDict
 from elk.extraction.extraction import extract
 from elk.files import create_output_directory, save_config, save_meta
+from elk.logging import save_debug_log
 from elk.training.preprocessing import normalize
-from elk.utils.csv import Log, write_iterator_to_file
 from elk.utils.data_utils import get_layers, select_train_val_splits
 from elk.utils.typing import assert_type, int16_to_float32
 
@@ -94,10 +95,8 @@ def concatenate(self, layers):
 
  def apply_to_layers(
  self,
- func: Callable[[int], Log],
+ func: Callable[[int], pd.Series],
  num_devices: int,
- to_csv_line: Callable[[Log], list[str]],
- csv_columns: list[str],
  ):
  """Apply a function to each layer of the dataset in parallel
  and writes the results to a CSV file.
@@ -120,15 +119,14 @@ def apply_to_layers(
  # Should we write to different CSV files for elicit vs eval?
  with mp.Pool(num_devices) as pool, open(self.out_dir / "eval.csv", "w") as f:
  mapper = pool.imap_unordered if num_devices > 1 else map
- iterator: Iterator[Log] = tqdm( # type: ignore
- mapper(func, layers), total=len(layers)
- )
- write_iterator_to_file(
- iterator=iterator,
- file=f,
- debug=self.cfg.debug,
- dataset=self.dataset,
- out_dir=self.out_dir,
- csv_columns=csv_columns,
- to_csv_line=to_csv_line,
- )
+ row_buf = []
+
+ try:
+ for row in tqdm(mapper(func, layers), total=len(layers)):
+ row_buf.append(row)
+ finally:
+ # Make sure the CSV is written even if we crash or get interrupted
+ df = pd.DataFrame(row_buf).sort_values(by="layer")
+ df.to_csv(f, index=False)
+ if self.cfg.debug:
+ save_debug_log(self.dataset, self.out_dir)

elk/training/baseline.py

@@ -1,6 +1,6 @@
 import pickle
 from pathlib import Path
-from typing import Tuple
+from typing import NamedTuple, Tuple
 
 import torch
 from sklearn.metrics import accuracy_score, roc_auc_score

elk/training/reporter.py

@@ -1,16 +1,18 @@
 """An ELK reporter network."""
 
-from ..calibration import CalibrationError
-from .classifier import Classifier
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from pathlib import Path
-from simple_parsing.helpers import Serializable
-from sklearn.metrics import roc_auc_score
-from torch import Tensor
 from typing import Literal, NamedTuple, Optional, Union
+
 import torch
 import torch.nn as nn
+from simple_parsing.helpers import Serializable
+from sklearn.metrics import roc_auc_score
+from torch import Tensor
+
+from ..calibration import CalibrationError
+from .classifier import Classifier
 
 
 class EvalResult(NamedTuple):