Upload 16 files

.gitattributes

@@ -33,3 +33,9 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+data/cifar-10-batches-py/data_batch_1 filter=lfs diff=lfs merge=lfs -text
+data/cifar-10-batches-py/data_batch_2 filter=lfs diff=lfs merge=lfs -text
+data/cifar-10-batches-py/data_batch_3 filter=lfs diff=lfs merge=lfs -text
+data/cifar-10-batches-py/data_batch_4 filter=lfs diff=lfs merge=lfs -text
+data/cifar-10-batches-py/data_batch_5 filter=lfs diff=lfs merge=lfs -text
+data/cifar-10-batches-py/test_batch filter=lfs diff=lfs merge=lfs -text

checkpoints/detector.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4f1e2a8a72bad1870028531540d8973cc071e74ce7976992cfd033db03db300a
+size 2116641

data/cifar-10-batches-py/data_batch_1

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:54636561a3ce25bd3e19253c6b0d8538147b0ae398331ac4a2d86c6d987368cd
+size 31035704

data/cifar-10-batches-py/data_batch_2

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:766b2cef9fbc745cf056b3152224f7cf77163b330ea9a15f9392beb8b89bc5a8
+size 31035320

data/cifar-10-batches-py/data_batch_3

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0f00d98ebfb30b3ec0ad19f9756dc2630b89003e10525f5e148445e82aa6a1f9
+size 31035999

data/cifar-10-batches-py/data_batch_4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3f7bb240661948b8f4d53e36ec720d8306f5668bd0071dcb4e6c947f78e9682b
+size 31035696

data/cifar-10-batches-py/data_batch_5

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d91802434d8376bbaeeadf58a737e3a1b12ac839077e931237e0dcd43adcb154
+size 31035623

data/cifar-10-batches-py/readme.html

@@ -0,0 +1 @@
+<meta HTTP-EQUIV="REFRESH" content="0; url=http://www.cs.toronto.edu/~kriz/cifar.html">

data/cifar-10-batches-py/test_batch

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f53d8d457504f7cff4ea9e021afcf0e0ad8e24a91f3fc42091b8adef61157831
+size 31035526

data/cifar-10-python.tar.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6d958be074577803d12ecdefd02955f39262c83c16fe9348329d7fe0b5c001ce
+size 170498071

detector.py

@@ -0,0 +1,230 @@
+"""
+Detector — "імунна клітина" системи.
+Окрема модель, що розрізняє чисті vs отруєні зразки.
+
+Архітектура:
+    - Encoder (CNN) → embedding простір
+    - Poison head: бінарна класифікація clean/poisoned
+    - Attack-type head: класифікація типу атаки (для аналізу)
+    - Memory bank: зберігає embedding'и відомих атак ("лімфоцити пам'яті")
+
+Тренування:
+    - Класифікаційний loss (cross-entropy)
+    - Contrastive loss (SupCon) — чисті embedding'и кластеризуються разом
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Dict, Tuple, Optional
+
+
+class DetectorEncoder(nn.Module):
+    """Малий CNN-енкодер. Вистачає для CIFAR-розмірів."""
+
+    def __init__(self, in_channels: int = 3, embed_dim: int = 128):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Conv2d(in_channels, 32, 3, padding=1),
+            nn.BatchNorm2d(32),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2),  # 32 → 16
+            nn.Conv2d(32, 64, 3, padding=1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2),  # 16 → 8
+            nn.Conv2d(64, 128, 3, padding=1),
+            nn.BatchNorm2d(128),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2),  # 8 → 4
+            nn.Conv2d(128, 128, 3, padding=1),
+            nn.BatchNorm2d(128),
+            nn.ReLU(inplace=True),
+            nn.AdaptiveAvgPool2d(1),
+            nn.Flatten(),
+        )
+        self.projection = nn.Linear(128, embed_dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        h = self.net(x)
+        return self.projection(h)
+
+
+class Detector(nn.Module):
+    """
+    Детектор з двома головами + memory bank.
+
+    Forward повертає:
+        embeddings: (B, embed_dim) — нормалізовані embedding'и
+        poison_logits: (B, 2) — clean vs poisoned
+        attack_logits: (B, num_attack_types) — який тип атаки
+    """
+
+    def __init__(
+        self,
+        in_channels: int = 3,
+        embed_dim: int = 128,
+        num_attack_types: int = 5,  # clean + 4 attack types
+    ):
+        super().__init__()
+        self.encoder = DetectorEncoder(in_channels, embed_dim)
+        self.poison_head = nn.Linear(embed_dim, 2)
+        self.attack_head = nn.Linear(embed_dim, num_attack_types)
+        self.embed_dim = embed_dim
+
+        # Memory bank — зберігаємо embedding'и відомих атак
+        self.register_buffer("memory_embeds", torch.zeros(0, embed_dim))
+        self.register_buffer("memory_attack_ids", torch.zeros(0, dtype=torch.long))
+
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        embed = self.encoder(x)
+        embed_normalized = F.normalize(embed, dim=-1)
+        poison_logits = self.poison_head(embed)
+        attack_logits = self.attack_head(embed)
+        return embed_normalized, poison_logits, attack_logits
+
+    @torch.no_grad()
+    def poison_probability(self, x: torch.Tensor) -> torch.Tensor:
+        """Повертає P(poisoned) для кожного зразка."""
+        self.eval()
+        _, poison_logits, _ = self.forward(x)
+        probs = F.softmax(poison_logits, dim=-1)
+        return probs[:, 1]
+
+    @torch.no_grad()
+    def trust_weights(self, x: torch.Tensor, soft: bool = True) -> torch.Tensor:
+        """
+        Ваги довіри для loss-зважування.
+        soft=True: ваги в [0, 1] = 1 - P(poisoned)
+        soft=False: hard rejection — 1.0 якщо clean, 0.0 якщо poisoned
+        """
+        p = self.poison_probability(x)
+        if soft:
+            return 1.0 - p
+        return (p < 0.5).float()
+
+    @torch.no_grad()
+    def update_memory(
+        self,
+        embeddings: torch.Tensor,
+        attack_ids: torch.Tensor,
+        max_size: int = 2000,
+    ):
+        """Додає нові сигнатури в memory bank ('лімфоцити пам'яті')."""
+        self.memory_embeds = torch.cat([self.memory_embeds, embeddings.detach()])
+        self.memory_attack_ids = torch.cat([self.memory_attack_ids, attack_ids])
+
+        if len(self.memory_embeds) > max_size:
+            # Залишаємо найсвіжіші
+            self.memory_embeds = self.memory_embeds[-max_size:]
+            self.memory_attack_ids = self.memory_attack_ids[-max_size:]
+
+    @torch.no_grad()
+    def memory_lookup(self, x: torch.Tensor, k: int = 5) -> Optional[torch.Tensor]:
+        """
+        Шукає k найближчих зразків у memory bank.
+        Повертає вектор оцінок схожості з відомими атаками.
+        """
+        if len(self.memory_embeds) == 0:
+            return None
+        embed, _, _ = self.forward(x)
+        # Косинусна схожість (embed вже нормалізований)
+        sim = torch.matmul(embed, self.memory_embeds.T)  # (B, M)
+        top_k_sim, top_k_idx = sim.topk(min(k, sim.size(1)), dim=1)
+        # Чи нагадує зразок щось з memory (середня top-k схожість)
+        return top_k_sim.mean(dim=1)
+
+
+def supervised_contrastive_loss(
+    embeddings: torch.Tensor, labels: torch.Tensor, temperature: float = 0.1
+) -> torch.Tensor:
+    """
+    SupCon loss (Khosla et al., 2020).
+    Тягне зразки однієї мітки ближче, штовхає інші геть.
+    Для нас: чисті embedding'и → один кластер, отруєні → окремо.
+    """
+    device = embeddings.device
+    batch_size = embeddings.size(0)
+
+    # similarity matrix
+    sim = torch.matmul(embeddings, embeddings.T) / temperature
+
+    # числова стабільність
+    sim_max, _ = sim.max(dim=1, keepdim=True)
+    sim = sim - sim_max.detach()
+
+    # маска "позитивних пар" (та сама мітка, виключаючи self)
+    labels = labels.contiguous().view(-1, 1)
+    pos_mask = torch.eq(labels, labels.T).float().to(device)
+    self_mask = torch.eye(batch_size, device=device)
+    pos_mask = pos_mask - self_mask  # виключити діагональ
+
+    # log_prob
+    exp_sim = torch.exp(sim) * (1 - self_mask)  # виключити self
+    log_prob = sim - torch.log(exp_sim.sum(dim=1, keepdim=True) + 1e-12)
+
+    # середнє по позитивах
+    pos_count = pos_mask.sum(dim=1)
+    pos_count = torch.clamp(pos_count, min=1.0)
+    mean_log_prob_pos = (pos_mask * log_prob).sum(dim=1) / pos_count
+
+    return -mean_log_prob_pos.mean()
+
+
+def detector_loss(
+    embeddings: torch.Tensor,
+    poison_logits: torch.Tensor,
+    attack_logits: torch.Tensor,
+    is_poisoned: torch.Tensor,
+    attack_ids: torch.Tensor,
+    use_contrastive: bool = True,
+    contrast_weight: float = 0.5,
+    attack_weight: float = 0.3,
+) -> Tuple[torch.Tensor, Dict[str, float]]:
+    """
+    Комбінований loss для тренування Detector'а.
+
+    Args:
+        embeddings: (B, D) нормалізовані
+        poison_logits: (B, 2)
+        attack_logits: (B, num_attack_types)
+        is_poisoned: (B,) bool
+        attack_ids: (B,) long — 0 для clean, 1..N для типів атак
+    """
+    is_poisoned_long = is_poisoned.long()
+
+    # 1. Бінарна класифікація clean/poisoned
+    cls_loss = F.cross_entropy(poison_logits, is_poisoned_long)
+
+    # 2. Multi-class класифікація типу атаки
+    atk_loss = F.cross_entropy(attack_logits, attack_ids)
+
+    total = cls_loss + attack_weight * atk_loss
+    metrics = {"cls": cls_loss.item(), "attack_cls": atk_loss.item()}
+
+    # 3. SupCon — для розділення в embedding-просторі
+    if use_contrastive and embeddings.size(0) > 2:
+        contrast = supervised_contrastive_loss(embeddings, is_poisoned_long)
+        total = total + contrast_weight * contrast
+        metrics["contrast"] = contrast.item()
+
+    metrics["total"] = total.item()
+    return total, metrics
+
+
+if __name__ == "__main__":
+    # Швидкий тест
+    detector = Detector(in_channels=3, embed_dim=128, num_attack_types=5)
+    x = torch.rand(8, 3, 32, 32)
+    embed, poison_logits, attack_logits = detector(x)
+
+    print(f"Embedding shape: {embed.shape}")
+    print(f"Poison logits shape: {poison_logits.shape}")
+    print(f"Attack logits shape: {attack_logits.shape}")
+    print(f"Trust weights: {detector.trust_weights(x).tolist()}")
+
+    # Тест loss'а
+    is_poisoned = torch.randint(0, 2, (8,)).bool()
+    attack_ids = torch.randint(0, 5, (8,))
+    loss, metrics = detector_loss(embed, poison_logits, attack_logits, is_poisoned, attack_ids)
+    print(f"Loss: {loss.item():.4f}, metrics: {metrics}")

hf_data.py

@@ -0,0 +1,185 @@
+"""
+HuggingFace datasets — обгортка для PyTorch DataLoader'а.
+
+Чому це краще ніж torchvision:
+    - Більше готових датасетів (CIFAR-100, Tiny ImageNet, Food-101, etc.)
+    - Уніфікований API
+    - Можна стрімити великі датасети (для майбутнього ImageNet)
+    - Легко підмінити датасет — тільки змінити рядок
+
+Решта пайплайну (poison_generator, detector, models) НЕ ЗМІНЮЄТЬСЯ:
+    DataLoader повертає звичайні (image_tensor, label) пари —
+    рівно як torchvision робив.
+"""
+
+from typing import Optional, Tuple
+import torch
+from torch.utils.data import Dataset, DataLoader
+from torchvision import transforms
+from datasets import load_dataset
+
+
+class HFImageDataset(Dataset):
+    """
+    Адаптер між HuggingFace Dataset і PyTorch Dataset.
+
+    HF повертає dict з PIL-зображенням, нам треба (tensor, label).
+    """
+
+    def __init__(
+        self,
+        hf_dataset,
+        image_col: str = "img",
+        label_col: str = "label",
+        transform: Optional[transforms.Compose] = None,
+    ):
+        self.ds = hf_dataset
+        self.image_col = image_col
+        self.label_col = label_col
+        # За замовч.: PIL → Tensor у діапазоні [0, 1]
+        # (важливо для poison_generator'а, бо clean_label_attack очікує [0,1])
+        self.transform = transform or transforms.ToTensor()
+
+    def __len__(self) -> int:
+        return len(self.ds)
+
+    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, int]:
+        item = self.ds[idx]
+        img = item[self.image_col]
+        label = item[self.label_col]
+
+        # MNIST у HF — grayscale L, CIFAR — RGB. transforms.ToTensor() справиться з обома.
+        if self.transform:
+            img = self.transform(img)
+
+        # Гарантуємо 3 канали для уніфікації (опціонально)
+        # if img.size(0) == 1:
+        #     img = img.repeat(3, 1, 1)
+
+        return img, label
+
+
+# Реєстр підтримуваних датасетів: hf_name → конфіг
+DATASET_CONFIGS = {
+    "cifar10": {
+        "hf_name": "cifar10",
+        "image_col": "img",
+        "label_col": "label",
+        "in_channels": 3,
+        "num_classes": 10,
+        "image_size": 32,
+    },
+    "cifar100": {
+        "hf_name": "cifar100",
+        "image_col": "img",
+        "label_col": "fine_label",  # CIFAR-100 має fine та coarse мітки
+        "in_channels": 3,
+        "num_classes": 100,
+        "image_size": 32,
+    },
+    "mnist": {
+        "hf_name": "mnist",
+        "image_col": "image",
+        "label_col": "label",
+        "in_channels": 1,
+        "num_classes": 10,
+        "image_size": 28,
+    },
+    "tiny_imagenet": {
+        "hf_name": "zh-plus/tiny-imagenet",
+        "image_col": "image",
+        "label_col": "label",
+        "in_channels": 3,
+        "num_classes": 200,
+        "image_size": 64,
+    },
+    "fashion_mnist": {
+        "hf_name": "fashion_mnist",
+        "image_col": "image",
+        "label_col": "label",
+        "in_channels": 1,
+        "num_classes": 10,
+        "image_size": 28,
+    },
+}
+
+
+def get_dataloaders(
+    dataset: str,
+    batch_size: int,
+    cache_dir: Optional[str] = None,
+    num_workers: int = 2,
+    resize_to: Optional[int] = None,
+) -> Tuple[DataLoader, DataLoader, int, int]:
+    """
+    Завантажує датасет з HuggingFace Hub і повертає PyTorch DataLoader'и.
+
+    Args:
+        dataset: ключ з DATASET_CONFIGS ("cifar10", "cifar100", тощо)
+        batch_size: розмір батча
+        cache_dir: куди кешувати датасет (None → за замовч. ~/.cache/huggingface)
+        num_workers: процесів для завантаження даних
+        resize_to: якщо вказано — змінює розмір зображень (для уніфікації архітектури)
+
+    Returns:
+        (train_loader, test_loader, in_channels, num_classes)
+    """
+    if dataset not in DATASET_CONFIGS:
+        raise ValueError(
+            f"Unknown dataset: {dataset}. Available: {list(DATASET_CONFIGS.keys())}"
+        )
+
+    cfg = DATASET_CONFIGS[dataset]
+
+    # Завантажуємо обидва спліти (HF датасети бувають з різними іменами для test)
+    print(f"Loading {cfg['hf_name']} from HuggingFace Hub...")
+    ds = load_dataset(cfg["hf_name"], cache_dir=cache_dir)
+
+    # HF датасети мають різні імена для test спліту
+    train_split = "train"
+    test_split = "test" if "test" in ds else ("validation" if "validation" in ds else "valid")
+
+    print(f"Splits available: {list(ds.keys())}, using train='{train_split}', test='{test_split}'")
+    print(f"Train size: {len(ds[train_split])}, Test size: {len(ds[test_split])}")
+
+    # Трансформація: PIL → Tensor [0,1], опційно resize
+    tx_list = []
+    if resize_to is not None:
+        tx_list.append(transforms.Resize((resize_to, resize_to)))
+    tx_list.append(transforms.ToTensor())
+    transform = transforms.Compose(tx_list)
+
+    train_set = HFImageDataset(
+        ds[train_split],
+        image_col=cfg["image_col"],
+        label_col=cfg["label_col"],
+        transform=transform,
+    )
+    test_set = HFImageDataset(
+        ds[test_split],
+        image_col=cfg["image_col"],
+        label_col=cfg["label_col"],
+        transform=transform,
+    )
+
+    train_loader = DataLoader(
+        train_set, batch_size=batch_size, shuffle=True,
+        num_workers=num_workers, pin_memory=torch.cuda.is_available(),
+    )
+    test_loader = DataLoader(
+        test_set, batch_size=batch_size, shuffle=False,
+        num_workers=num_workers, pin_memory=torch.cuda.is_available(),
+    )
+
+    return train_loader, test_loader, cfg["in_channels"], cfg["num_classes"]
+
+
+if __name__ == "__main__":
+    # Швидкий тест
+    train_loader, test_loader, in_ch, n_cls = get_dataloaders("cifar10", batch_size=32)
+    x, y = next(iter(train_loader))
+    print(f"\nBatch shape: {x.shape}")
+    print(f"Labels shape: {y.shape}")
+    print(f"Pixel range: [{x.min():.3f}, {x.max():.3f}]")
+    print(f"In channels: {in_ch}, Classes: {n_cls}")
+    print(f"Train batches: {len(train_loader)}, Test batches: {len(test_loader)}")

models.py

@@ -0,0 +1,102 @@
+"""
+Protected Model — основна модель, яку ми захищаємо від отруєння.
+
+Це невеликий ResNet для CIFAR-розмірів.
+Її задача — звичайна класифікація.
+Захист реалізований у train.py через зважування loss за trust_weights від Detector'а.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class BasicBlock(nn.Module):
+    """Стандартний ResNet block."""
+
+    expansion = 1
+
+    def __init__(self, in_channels: int, out_channels: int, stride: int = 1):
+        super().__init__()
+        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, stride, 1, bias=False)
+        self.bn1 = nn.BatchNorm2d(out_channels)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, 1, 1, bias=False)
+        self.bn2 = nn.BatchNorm2d(out_channels)
+
+        self.shortcut = nn.Sequential()
+        if stride != 1 or in_channels != out_channels:
+            self.shortcut = nn.Sequential(
+                nn.Conv2d(in_channels, out_channels, 1, stride, bias=False),
+                nn.BatchNorm2d(out_channels),
+            )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        out = F.relu(self.bn1(self.conv1(x)), inplace=True)
+        out = self.bn2(self.conv2(out))
+        out = out + self.shortcut(x)
+        return F.relu(out, inplace=True)
+
+
+class ProtectedModel(nn.Module):
+    """
+    Малий ResNet (~ResNet-14) для CIFAR.
+    """
+
+    def __init__(self, num_classes: int = 10, in_channels: int = 3):
+        super().__init__()
+
+        self.stem = nn.Sequential(
+            nn.Conv2d(in_channels, 64, 3, 1, 1, bias=False),
+            nn.BatchNorm2d(64),
+            nn.ReLU(inplace=True),
+        )
+
+        self.layer1 = self._make_layer(64, 64, num_blocks=2, stride=1)
+        self.layer2 = self._make_layer(64, 128, num_blocks=2, stride=2)
+        self.layer3 = self._make_layer(128, 256, num_blocks=2, stride=2)
+
+        self.pool = nn.AdaptiveAvgPool2d(1)
+        self.fc = nn.Linear(256, num_classes)
+
+    def _make_layer(self, in_ch: int, out_ch: int, num_blocks: int, stride: int) -> nn.Sequential:
+        layers = [BasicBlock(in_ch, out_ch, stride)]
+        for _ in range(num_blocks - 1):
+            layers.append(BasicBlock(out_ch, out_ch, 1))
+        return nn.Sequential(*layers)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.stem(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.pool(x).flatten(1)
+        return self.fc(x)
+
+
+def weighted_cross_entropy(
+    logits: torch.Tensor, targets: torch.Tensor, weights: torch.Tensor
+) -> torch.Tensor:
+    """
+    Cross-entropy з вагою на кожен зразок.
+    Це ключове місце "імунного захисту":
+        weights близько 1 → нормальне навчання
+        weights близько 0 → зразок майже не впливає
+    """
+    per_sample_loss = F.cross_entropy(logits, targets, reduction="none")
+    # Нормалізуємо за сумою ваг, щоб масштаб loss'у був стабільним
+    weight_sum = weights.sum().clamp(min=1e-6)
+    return (per_sample_loss * weights).sum() / weight_sum
+
+
+if __name__ == "__main__":
+    model = ProtectedModel(num_classes=10)
+    x = torch.rand(4, 3, 32, 32)
+    out = model(x)
+    print(f"Output shape: {out.shape}")
+    print(f"Params: {sum(p.numel() for p in model.parameters()) / 1e6:.2f}M")
+
+    # Тест weighted loss
+    targets = torch.randint(0, 10, (4,))
+    weights = torch.tensor([1.0, 0.1, 1.0, 0.0])  # 2-й трохи підозрілий, 4-й — точно отрута
+    loss = weighted_cross_entropy(out, targets, weights)
+    print(f"Weighted loss: {loss.item():.4f}")

poison_generator.py

@@ -0,0 +1,208 @@
+"""
+Poison Generator — "лабораторія патогенів".
+Створює різноманітні отруєні зразки для тренування Detector'а.
+
+Підтримувані атаки:
+    - LabelFlipAttack: зміна мітки класу
+    - BackdoorAttack: додавання тригер-патчу + зміна мітки
+    - CleanLabelAttack: непомітна пертурбація зі збереженням мітки
+    - FeatureCorruptionAttack: шум/артефакти
+"""
+
+import torch
+import numpy as np
+from typing import Tuple, List, Optional
+
+
+class BaseAttack:
+    """Базовий клас атаки. Кожна атака повертає (x_poisoned, y_poisoned)."""
+
+    name: str = "base"
+
+    def __call__(self, x: torch.Tensor, y: int) -> Tuple[torch.Tensor, int]:
+        raise NotImplementedError
+
+
+class LabelFlipAttack(BaseAttack):
+    """Перевертає мітку класу. Найпростіший тип отруєння."""
+
+    name = "label_flip"
+
+    def __init__(self, num_classes: int = 10, target: Optional[int] = None):
+        self.num_classes = num_classes
+        self.target = target  # якщо None — випадковий інший клас
+
+    def __call__(self, x: torch.Tensor, y: int) -> Tuple[torch.Tensor, int]:
+        if self.target is not None:
+            new_y = self.target
+        else:
+            # Випадковий клас, але не оригінальний
+            choices = [c for c in range(self.num_classes) if c != y]
+            new_y = int(np.random.choice(choices))
+        return x.clone(), new_y
+
+
+class BackdoorAttack(BaseAttack):
+    """
+    Додає тригер-патч (квадратик) у кут зображення + змінює мітку на target_class.
+    Класична backdoor / trojan-атака: модель починає асоціювати патерн з класом.
+    """
+
+    name = "backdoor"
+
+    def __init__(
+        self,
+        trigger_size: int = 4,
+        trigger_value: float = 1.0,
+        target_class: int = 0,
+        position: str = "bottom_right",
+    ):
+        self.trigger_size = trigger_size
+        self.trigger_value = trigger_value
+        self.target_class = target_class
+        self.position = position
+
+    def __call__(self, x: torch.Tensor, y: int) -> Tuple[torch.Tensor, int]:
+        x_p = x.clone()
+        s = self.trigger_size
+
+        if self.position == "bottom_right":
+            x_p[..., -s:, -s:] = self.trigger_value
+        elif self.position == "top_left":
+            x_p[..., :s, :s] = self.trigger_value
+        elif self.position == "center":
+            h, w = x_p.shape[-2:]
+            ch, cw = h // 2, w // 2
+            x_p[..., ch - s // 2 : ch + s // 2, cw - s // 2 : cw + s // 2] = self.trigger_value
+
+        return x_p, self.target_class
+
+
+class CleanLabelAttack(BaseAttack):
+    """
+    Clean-label attack: непомітна пертурбація, мітка не змінюється.
+    Використовує adversarial noise (Gaussian або PGD-style).
+    Найхитріша атака — її важко знайти просто переглянувши датасет.
+    """
+
+    name = "clean_label"
+
+    def __init__(self, epsilon: float = 0.03, mode: str = "gaussian"):
+        self.epsilon = epsilon
+        self.mode = mode  # "gaussian" або "fgsm"
+
+    def __call__(self, x: torch.Tensor, y: int) -> Tuple[torch.Tensor, int]:
+        if self.mode == "gaussian":
+            noise = torch.randn_like(x) * self.epsilon
+        else:
+            # Простий sign-noise (наближення FGSM без градієнтів моделі)
+            noise = torch.sign(torch.randn_like(x)) * self.epsilon
+
+        x_p = torch.clamp(x + noise, 0.0, 1.0)
+        return x_p, y  # мітка залишається!
+
+
+class FeatureCorruptionAttack(BaseAttack):
+    """
+    Корумпує частину пікселів — імітує зіпсовані дані в датасеті.
+    Може змінювати або не змінювати мітку.
+    """
+
+    name = "feature_corruption"
+
+    def __init__(self, corruption_ratio: float = 0.2, flip_label: bool = False, num_classes: int = 10):
+        self.corruption_ratio = corruption_ratio
+        self.flip_label = flip_label
+        self.num_classes = num_classes
+
+    def __call__(self, x: torch.Tensor, y: int) -> Tuple[torch.Tensor, int]:
+        x_p = x.clone()
+        mask = torch.rand_like(x_p) < self.corruption_ratio
+        random_pixels = torch.rand_like(x_p)
+        x_p[mask] = random_pixels[mask]
+
+        if self.flip_label:
+            new_y = int(np.random.choice([c for c in range(self.num_classes) if c != y]))
+            return x_p, new_y
+        return x_p, y
+
+
+class PoisonGenerator:
+    """
+    Оркеструє різні атаки. Отруює певний відсоток батчу.
+
+    Повертає:
+        - x_batch: отруєний батч
+        - y_batch: мітки (можливо змінені)
+        - is_poisoned: bool-маска, які зразки отруєні
+        - attack_types: список рядків з типом атаки (для аналізу)
+    """
+
+    def __init__(
+        self,
+        attacks: Optional[List[BaseAttack]] = None,
+        poison_ratio: float = 0.3,
+        num_classes: int = 10,
+    ):
+        if attacks is None:
+            attacks = [
+                LabelFlipAttack(num_classes=num_classes),
+                BackdoorAttack(trigger_size=4, target_class=0),
+                CleanLabelAttack(epsilon=0.05),
+                FeatureCorruptionAttack(corruption_ratio=0.2),
+            ]
+        self.attacks = attacks
+        self.poison_ratio = poison_ratio
+        self.num_classes = num_classes
+        self.attack_names = [a.name for a in attacks]
+
+    def poison_batch(
+        self,
+        x_batch: torch.Tensor,
+        y_batch: torch.Tensor,
+        poison_ratio: Optional[float] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, List[str]]:
+        batch_size = x_batch.size(0)
+        new_x = x_batch.clone()
+        new_y = y_batch.clone()
+        is_poisoned = torch.zeros(batch_size, dtype=torch.bool)
+        attack_types = ["clean"] * batch_size
+        ratio = poison_ratio if poison_ratio is not None else self.poison_ratio
+
+        for i in range(batch_size):
+            if np.random.rand() < ratio:
+                attack = np.random.choice(self.attacks)
+                x_p, y_p = attack(x_batch[i], int(y_batch[i].item()))
+                new_x[i] = x_p
+                new_y[i] = y_p
+                is_poisoned[i] = True
+                attack_types[i] = attack.name
+
+        return new_x, new_y, is_poisoned, attack_types
+
+    def poison_all(
+        self, x_batch: torch.Tensor, y_batch: torch.Tensor, attack: BaseAttack
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Отруює весь батч однією конкретною атакою — для тестування."""
+        new_x = x_batch.clone()
+        new_y = y_batch.clone()
+        for i in range(x_batch.size(0)):
+            x_p, y_p = attack(x_batch[i], int(y_batch[i].item()))
+            new_x[i] = x_p
+            new_y[i] = y_p
+        return new_x, new_y
+
+
+if __name__ == "__main__":
+    # Швидкий тест
+    x = torch.rand(8, 3, 32, 32)
+    y = torch.randint(0, 10, (8,))
+
+    gen = PoisonGenerator(poison_ratio=0.5)
+    new_x, new_y, mask, types = gen.poison_batch(x, y)
+
+    print(f"Original shape: {x.shape}")
+    print(f"Poisoned mask: {mask.tolist()}")
+    print(f"Attack types: {types}")
+    print(f"Label changes: {(new_y != y).tolist()}")
+    print(f"Pixel diff (max): {(new_x - x).abs().max().item():.4f}")

requirements.txt

@@ -0,0 +1,3 @@
+torch>=2.0.0
+torchvision>=0.15.0
+numpy>=1.24.0

train.py

@@ -0,0 +1,358 @@
+"""
+Повний пайплайн навчання з захистом від отруєння.
+
+Процес:
+    Фаза 1: Тренуємо Detector на згенерованих clean/poisoned парах.
+    Фаза 2: Тренуємо дві моделі:
+        a) Baseline — без захисту, на отруєному датасеті
+        b) Protected — з Detector-зважуванням loss
+    Фаза 3: Оцінюємо обидві на:
+        - Clean test set (звичайна точність)
+        - Backdoored test set (наскільки атака спрацьовує — менше = краще)
+
+Запуск:
+    python train.py --dataset cifar10 --epochs 20 --poison_ratio 0.3
+"""
+
+import argparse
+import os
+import time
+from typing import Dict, Tuple
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader
+
+from poison_generator import (
+    PoisonGenerator,
+    LabelFlipAttack,
+    BackdoorAttack,
+    CleanLabelAttack,
+    FeatureCorruptionAttack,
+)
+from detector import Detector, detector_loss
+from models import ProtectedModel, weighted_cross_entropy
+from hf_data import get_dataloaders, DATASET_CONFIGS
+
+
+# ---- Мапінг назв атак на ID (для attack_head) ----
+ATTACK_NAME_TO_ID = {
+    "clean": 0,
+    "label_flip": 1,
+    "backdoor": 2,
+    "clean_label": 3,
+    "feature_corruption": 4,
+}
+NUM_ATTACK_TYPES = len(ATTACK_NAME_TO_ID)
+
+
+def make_attack_id_tensor(attack_types: list) -> torch.Tensor:
+    return torch.tensor([ATTACK_NAME_TO_ID[t] for t in attack_types], dtype=torch.long)
+
+
+# =============================================================================
+# ФАЗА 1: ТРЕНУВАННЯ DETECTOR'А
+# =============================================================================
+def train_detector(
+    detector: Detector,
+    train_loader: DataLoader,
+    poison_gen: PoisonGenerator,
+    epochs: int,
+    device: torch.device,
+    lr: float = 1e-3,
+):
+    print("\n" + "=" * 60)
+    print("ФАЗА 1: Тренування Detector'а")
+    print("=" * 60)
+
+    optimizer = optim.Adam(detector.parameters(), lr=lr)
+    detector.train()
+
+    for epoch in range(epochs):
+        epoch_loss = 0.0
+        correct = 0
+        total = 0
+        attack_correct = 0
+        start = time.time()
+
+        for batch_idx, (x, y) in enumerate(train_loader):
+            x, y = x.to(device), y.to(device)
+
+            # Генеруємо отруєний батч (60% отрути для збалансованості при навчанні детектора)
+            x_mixed, _, is_poisoned, attack_types = poison_gen.poison_batch(x, y, poison_ratio=0.6)
+            x_mixed = x_mixed.to(device)
+            is_poisoned = is_poisoned.to(device)
+            attack_ids = make_attack_id_tensor(attack_types).to(device)
+
+            optimizer.zero_grad()
+            embed, poison_logits, attack_logits = detector(x_mixed)
+
+            loss, metrics = detector_loss(
+                embed, poison_logits, attack_logits,
+                is_poisoned, attack_ids,
+                use_contrastive=True,
+            )
+            loss.backward()
+            optimizer.step()
+
+            epoch_loss += metrics["total"]
+            preds = poison_logits.argmax(dim=-1)
+            correct += (preds == is_poisoned.long()).sum().item()
+            attack_correct += (attack_logits.argmax(dim=-1) == attack_ids).sum().item()
+            total += x.size(0)
+
+            # Оновлюємо memory bank на отруєних зразках
+            if is_poisoned.any():
+                detector.update_memory(
+                    embed[is_poisoned].detach(),
+                    attack_ids[is_poisoned],
+                )
+
+        elapsed = time.time() - start
+        print(
+            f"Detector epoch {epoch + 1}/{epochs} | "
+            f"loss={epoch_loss / len(train_loader):.4f} | "
+            f"binary_acc={100 * correct / total:.2f}% | "
+            f"attack_type_acc={100 * attack_correct / total:.2f}% | "
+            f"time={elapsed:.1f}s | "
+            f"memory_size={len(detector.memory_embeds)}"
+        )
+
+
+# =============================================================================
+# ФАЗА 2: ТРЕНУВАННЯ PROTECTED ТА BASELINE МОДЕЛЕЙ
+# =============================================================================
+def train_classifier(
+    model: nn.Module,
+    detector: Detector,  # None для baseline
+    train_loader: DataLoader,
+    poison_gen: PoisonGenerator,
+    epochs: int,
+    device: torch.device,
+    use_defense: bool,
+    lr: float = 1e-3,
+    name: str = "Model",
+):
+    print("\n" + "=" * 60)
+    print(f"ФАЗА 2: Тренування {name} (захист={'ON' if use_defense else 'OFF'})")
+    print("=" * 60)
+
+    optimizer = optim.Adam(model.parameters(), lr=lr)
+    if detector is not None:
+        detector.eval()
+
+    for epoch in range(epochs):
+        epoch_loss = 0.0
+        correct = 0
+        total = 0
+        avg_trust_poisoned = 0.0
+        avg_trust_clean = 0.0
+        num_poisoned = 0
+        num_clean = 0
+        start = time.time()
+
+        model.train()
+        for batch_idx, (x, y) in enumerate(train_loader):
+            x, y = x.to(device), y.to(device)
+
+            # Отруюємо частину датасету (це симуляція компрометованих даних)
+            x_p, y_p, is_poisoned, _ = poison_gen.poison_batch(x, y)
+            x_p = x_p.to(device)
+            y_p = y_p.to(device)
+            is_poisoned = is_poisoned.to(device)
+
+            optimizer.zero_grad()
+            logits = model(x_p)
+
+            if use_defense and detector is not None:
+                # Імунний захист: ваги довіри від детектора
+                trust = detector.trust_weights(x_p, soft=True)
+                loss = weighted_cross_entropy(logits, y_p, trust)
+
+                if is_poisoned.any():
+                    avg_trust_poisoned += trust[is_poisoned].sum().item()
+                    num_poisoned += is_poisoned.sum().item()
+                if (~is_poisoned).any():
+                    avg_trust_clean += trust[~is_poisoned].sum().item()
+                    num_clean += (~is_poisoned).sum().item()
+            else:
+                # Baseline — без захисту
+                loss = nn.functional.cross_entropy(logits, y_p)
+
+            loss.backward()
+            optimizer.step()
+
+            epoch_loss += loss.item()
+            correct += (logits.argmax(dim=-1) == y_p).sum().item()
+            total += x.size(0)
+
+        elapsed = time.time() - start
+        info = (
+            f"{name} epoch {epoch + 1}/{epochs} | "
+            f"loss={epoch_loss / len(train_loader):.4f} | "
+            f"train_acc(poisoned)={100 * correct / total:.2f}% | "
+            f"time={elapsed:.1f}s"
+        )
+        if use_defense and num_poisoned > 0:
+            info += (
+                f" | avg_trust(poisoned)={avg_trust_poisoned / num_poisoned:.3f}"
+                f" | avg_trust(clean)={avg_trust_clean / max(num_clean, 1):.3f}"
+            )
+        print(info)
+
+
+# =============================================================================
+# ФАЗА 3: ОЦІНКА
+# =============================================================================
+@torch.no_grad()
+def evaluate_clean(model: nn.Module, test_loader: DataLoader, device: torch.device) -> float:
+    model.eval()
+    correct = 0
+    total = 0
+    for x, y in test_loader:
+        x, y = x.to(device), y.to(device)
+        logits = model(x)
+        correct += (logits.argmax(dim=-1) == y).sum().item()
+        total += x.size(0)
+    return 100.0 * correct / total
+
+
+@torch.no_grad()
+def evaluate_backdoor(
+    model: nn.Module,
+    test_loader: DataLoader,
+    backdoor_attack: BackdoorAttack,
+    device: torch.device,
+) -> float:
+    """
+    Attack Success Rate (ASR): який % НЕ target-class зразків модель класифікує як target після додавання тригера.
+    МЕНШЕ = КРАЩЕ. Якщо захист працює — модель не повинна реагувати на тригер.
+    """
+    model.eval()
+    target = backdoor_attack.target_class
+    success = 0
+    total = 0
+
+    for x, y in test_loader:
+        # Пропускаємо зразки, які вже належать target класу
+        mask = y != target
+        if mask.sum() == 0:
+            continue
+
+        x_filtered = x[mask]
+        # Додаємо тригер
+        x_triggered = x_filtered.clone()
+        for i in range(x_triggered.size(0)):
+            x_triggered[i], _ = backdoor_attack(x_triggered[i], int(y[mask][i].item()))
+
+        x_triggered = x_triggered.to(device)
+        preds = model(x_triggered).argmax(dim=-1)
+        success += (preds == target).sum().item()
+        total += x_triggered.size(0)
+
+    return 100.0 * success / max(total, 1)
+
+
+# =============================================================================
+# MAIN
+# =============================================================================
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--dataset", type=str, default="cifar10",
+        choices=list(DATASET_CONFIGS.keys()),
+        help="HF dataset: cifar10, cifar100, mnist, tiny_imagenet, fashion_mnist",
+    )
+    parser.add_argument("--epochs_detector", type=int, default=5)
+    parser.add_argument("--epochs_classifier", type=int, default=10)
+    parser.add_argument("--batch_size", type=int, default=128)
+    parser.add_argument("--poison_ratio", type=float, default=0.3)
+    parser.add_argument("--lr", type=float, default=1e-3)
+    parser.add_argument("--cache_dir", type=str, default=None, help="HF cache dir (default: ~/.cache/huggingface)")
+    parser.add_argument("--save_dir", type=str, default="./checkpoints")
+    parser.add_argument("--num_workers", type=int, default=2)
+    parser.add_argument("--resize_to", type=int, default=None, help="Resize images to NxN (optional)")
+    parser.add_argument("--seed", type=int, default=42)
+    args = parser.parse_args()
+
+    torch.manual_seed(args.seed)
+    os.makedirs(args.save_dir, exist_ok=True)
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Device: {device}")
+    print(f"Args: {vars(args)}")
+
+    # Data — тепер через HuggingFace datasets
+    train_loader, test_loader, in_channels, num_classes = get_dataloaders(
+        args.dataset,
+        batch_size=args.batch_size,
+        cache_dir=args.cache_dir,
+        num_workers=args.num_workers,
+        resize_to=args.resize_to,
+    )
+
+    # Poison generator з усіма типами атак
+    target_class = 0  # бекдор завжди веде до класу 0
+    backdoor = BackdoorAttack(trigger_size=4, trigger_value=1.0, target_class=target_class)
+    poison_gen = PoisonGenerator(
+        attacks=[
+            LabelFlipAttack(num_classes=num_classes),
+            backdoor,
+            CleanLabelAttack(epsilon=0.05),
+            FeatureCorruptionAttack(corruption_ratio=0.2, num_classes=num_classes),
+        ],
+        poison_ratio=args.poison_ratio,
+        num_classes=num_classes,
+    )
+
+    # --- ФАЗА 1: Detector ---
+    detector = Detector(
+        in_channels=in_channels, embed_dim=128, num_attack_types=NUM_ATTACK_TYPES
+    ).to(device)
+    train_detector(detector, train_loader, poison_gen, args.epochs_detector, device, args.lr)
+    torch.save(detector.state_dict(), os.path.join(args.save_dir, "detector.pt"))
+
+    # --- ФАЗА 2: Baseline (без захисту) ---
+    baseline = ProtectedModel(num_classes=num_classes, in_channels=in_channels).to(device)
+    train_classifier(
+        baseline, None, train_loader, poison_gen,
+        args.epochs_classifier, device, use_defense=False, lr=args.lr, name="BASELINE"
+    )
+    torch.save(baseline.state_dict(), os.path.join(args.save_dir, "baseline.pt"))
+
+    # --- ФАЗА 2: Protected (з імунним захистом) ---
+    protected = ProtectedModel(num_classes=num_classes, in_channels=in_channels).to(device)
+    train_classifier(
+        protected, detector, train_loader, poison_gen,
+        args.epochs_classifier, device, use_defense=True, lr=args.lr, name="PROTECTED"
+    )
+    torch.save(protected.state_dict(), os.path.join(args.save_dir, "protected.pt"))
+
+    # --- ФАЗА 3: ОЦІНКА ---
+    print("\n" + "=" * 60)
+    print("ФАЗА 3: Фінальна оцінка")
+    print("=" * 60)
+
+    baseline_clean = evaluate_clean(baseline, test_loader, device)
+    protected_clean = evaluate_clean(protected, test_loader, device)
+    baseline_asr = evaluate_backdoor(baseline, test_loader, backdoor, device)
+    protected_asr = evaluate_backdoor(protected, test_loader, backdoor, device)
+
+    print(f"\n{'Метрика':<35} {'Baseline':<15} {'Protected':<15}")
+    print("-" * 65)
+    print(f"{'Clean accuracy ↑':<35} {baseline_clean:<15.2f} {protected_clean:<15.2f}")
+    print(f"{'Backdoor ASR ↓ (атака успішна %)':<35} {baseline_asr:<15.2f} {protected_asr:<15.2f}")
+
+    print("\nІнтерпретація:")
+    print(f"  • Clean accuracy: вища = краще (нормальна продуктивність)")
+    print(f"  • Backdoor ASR: нижча = краще (атака менш ефективна)")
+    if protected_asr < baseline_asr:
+        delta = baseline_asr - protected_asr
+        print(f"  ✓ Захист знизив успішність атаки на {delta:.2f} процентних пунктів!")
+    else:
+        print(f"  ✗ Захист не зменшив атаку — треба тюнити Detector.")
+
+
+if __name__ == "__main__":
+    main()