autoattack/state.py

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()