|
import json |
|
from dataclasses import dataclass, field, asdict |
|
from datetime import datetime |
|
from pathlib import Path |
|
from typing import Optional, Set |
|
import warnings |
|
|
|
import torch |
|
|
|
|
|
@dataclass |
|
class EvaluationState: |
|
_attacks_to_run: Set[str] |
|
path: Optional[Path] = None |
|
_run_attacks: Set[str] = field(default_factory=set) |
|
_robust_flags: Optional[torch.Tensor] = None |
|
_last_saved: datetime = datetime(1, 1, 1) |
|
_SAVE_TIMEOUT: int = 60 |
|
_clean_accuracy: float = float("nan") |
|
|
|
def to_disk(self, force: bool = False) -> None: |
|
seconds_since_last_save = (datetime.now() - |
|
self._last_saved).total_seconds() |
|
if self.path is None or (seconds_since_last_save < self._SAVE_TIMEOUT |
|
and not force): |
|
return |
|
self._last_saved = datetime.now() |
|
d = asdict(self) |
|
if self.robust_flags is not None: |
|
d["_robust_flags"] = d["_robust_flags"].cpu().tolist() |
|
d["_run_attacks"] = list(self._run_attacks) |
|
with self.path.open("w", ) as f: |
|
json.dump(d, f, default=str) |
|
|
|
@classmethod |
|
def from_disk(cls, path: Path) -> "EvaluationState": |
|
with path.open("r") as f: |
|
d = json.load(f) |
|
d["_robust_flags"] = torch.tensor(d["_robust_flags"], dtype=torch.bool) |
|
d["path"] = Path(d["path"]) |
|
if path != d["path"]: |
|
warnings.warn( |
|
UserWarning( |
|
"The given path is different from the one found in the state file." |
|
)) |
|
d["_last_saved"] = datetime.fromisoformat(d["_last_saved"]) |
|
return cls(**d) |
|
|
|
@property |
|
def robust_flags(self) -> Optional[torch.Tensor]: |
|
return self._robust_flags |
|
|
|
@robust_flags.setter |
|
def robust_flags(self, robust_flags: torch.Tensor) -> None: |
|
self._robust_flags = robust_flags |
|
self.to_disk(force=True) |
|
|
|
@property |
|
def run_attacks(self) -> Set[str]: |
|
return self._run_attacks |
|
|
|
def add_run_attack(self, attack: str) -> None: |
|
self._run_attacks.add(attack) |
|
self.to_disk() |
|
|
|
@property |
|
def attacks_to_run(self) -> Set[str]: |
|
return self._attacks_to_run |
|
|
|
@attacks_to_run.setter |
|
def attacks_to_run(self, _: Set[str]) -> None: |
|
raise ValueError("attacks_to_run cannot be set outside of the constructor") |
|
|
|
@property |
|
def clean_accuracy(self) -> float: |
|
return self._clean_accuracy |
|
|
|
@clean_accuracy.setter |
|
def clean_accuracy(self, accuracy) -> None: |
|
self._clean_accuracy = accuracy |
|
self.to_disk(force=True) |
|
|
|
@property |
|
def robust_accuracy(self) -> float: |
|
if self.robust_flags is None: |
|
raise ValueError("robust_flags is not set yet. Start the attack first.") |
|
if self.attacks_to_run - self.run_attacks: |
|
warnings.warn("You are checking `robust_accuracy` before all the attacks" |
|
" have been run.") |
|
return self.robust_flags.float().mean().item() |
