File size: 6,088 Bytes
0880420 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 | """User-facing zero-shot inference API for Time-RCD."""
from __future__ import annotations
from copy import deepcopy
from pathlib import Path
from typing import Literal, Optional, Tuple, Union
import numpy as np
HF_REPO_ID = "thu-sail-lab/Time-RCD"
CHECKPOINT_FILES = {
"uni": "best_model/pretrain_checkpoint_best_uni.pth",
"multi": "best_model/pretrain_checkpoint_best_multi.pth",
}
DEFAULT_WIN_SIZE = 5000
DEFAULT_BATCH_SIZE = {"uni": 64, "multi": 1}
class TimeRCDDetector:
"""Zero-shot time series anomaly detector powered by Time-RCD.
Parameters
----------
checkpoint_path:
Path to a local ``.pth`` checkpoint.
variant:
``"uni"`` for univariate series, ``"multi"`` for multivariate series.
win_size:
Sliding window length. Sequences shorter than this value use the full
sequence length instead.
batch_size:
Inference batch size. Defaults to 64 (uni) or 1 (multi).
device:
PyTorch device string, e.g. ``"cuda"`` or ``"cpu"``. Auto-detected when
omitted.
Notes
-----
The model is initialized on the first call to :meth:`predict`, when the
number of input features is known.
"""
def __init__(
self,
checkpoint_path: str,
variant: Literal["uni", "multi"] = "uni",
win_size: int = DEFAULT_WIN_SIZE,
batch_size: Optional[int] = None,
device: Optional[str] = None,
) -> None:
if variant not in CHECKPOINT_FILES:
raise ValueError(f"variant must be one of {list(CHECKPOINT_FILES)}, got {variant!r}")
self.variant = variant
self.win_size = win_size
self.batch_size = batch_size if batch_size is not None else DEFAULT_BATCH_SIZE[variant]
self.checkpoint_path = str(checkpoint_path)
self.device = device
self._tester = None
self._num_features: Optional[int] = None
def _ensure_tester(self, num_features: int) -> None:
"""Initialize a model compatible with ``num_features`` when needed."""
if self._tester is not None and self._num_features == num_features:
return
from ._core.time_rcd_config import default_config
from ._inference import TimeRCDPretrainTester
# ``default_config`` is a module-level template. Each detector must own
# its configuration so its runtime options cannot affect other instances.
config = deepcopy(default_config)
config.ts_config.patch_size = 16
config.win_size = self.win_size
config.batch_size = self.batch_size
config.ts_config.num_features = num_features
self._tester = TimeRCDPretrainTester(self.checkpoint_path, config)
self._num_features = num_features
if self.device is not None:
import torch
self._tester.device = torch.device(self.device)
self._tester.model.to(self._tester.device)
@classmethod
def from_pretrained(
cls,
repo_id: str = HF_REPO_ID,
variant: Literal["uni", "multi"] = "uni",
cache_dir: Optional[str] = None,
local_files_only: bool = False,
**kwargs,
) -> "TimeRCDDetector":
"""Load a checkpoint from Hugging Face Hub (cached locally after first use)."""
from huggingface_hub import hf_hub_download
checkpoint_path = hf_hub_download(
repo_id=repo_id,
filename=CHECKPOINT_FILES[variant],
cache_dir=cache_dir,
local_files_only=local_files_only,
)
return cls(checkpoint_path=checkpoint_path, variant=variant, **kwargs)
@classmethod
def from_local(
cls,
checkpoint_path: Union[str, Path],
variant: Literal["uni", "multi"] = "uni",
**kwargs,
) -> "TimeRCDDetector":
"""Load a checkpoint from a local path."""
path = Path(checkpoint_path)
if not path.is_file():
raise FileNotFoundError(f"Checkpoint not found: {path}")
return cls(checkpoint_path=str(path), variant=variant, **kwargs)
def predict(
self,
data: Union[np.ndarray, list],
return_logits: bool = False,
) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
"""Run zero-shot anomaly scoring on a single time series.
Parameters
----------
data:
Array of shape ``(T,)`` or ``(T, C)`` where ``T`` is time steps and
``C`` is the number of channels.
return_logits:
When ``True``, also return raw anomaly logits.
Returns
-------
scores:
Anomaly scores in ``[0, 1]``, one value per time step.
logits:
Returned only when ``return_logits=True``.
"""
data = np.asarray(data, dtype=np.float64)
if data.ndim == 1:
data = data.reshape(-1, 1)
elif data.ndim != 2:
raise ValueError(f"data must be 1D or 2D, got shape {data.shape}")
if self.variant == "uni" and data.shape[1] != 1:
raise ValueError(
"variant='uni' expects a univariate series with shape (T,) or (T, 1). "
f"Got {data.shape[1]} channels; use variant='multi' instead."
)
if self.variant == "multi" and data.shape[1] < 2:
raise ValueError(
"variant='multi' expects at least two channels with shape (T, C), where C > 1. "
"Use variant='uni' for a univariate series."
)
original_length = data.shape[0]
self._ensure_tester(data.shape[1])
assert self._tester is not None
score_chunks, logit_chunks = self._tester.zero_shot(data)
scores = np.concatenate([np.asarray(chunk).reshape(-1) for chunk in score_chunks], axis=0)
logits = np.concatenate([np.asarray(chunk).reshape(-1) for chunk in logit_chunks], axis=0)
scores = scores[:original_length]
logits = logits[:original_length]
if return_logits:
return scores, logits
return scores
|