File size: 5,409 Bytes
998bb30 |
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 |
import numpy as np
import torch
from torch.utils.data import Dataset
from anti_kd_backdoor.data.dataset.base import (
DatasetInterface,
IndexDataset,
IndexRatioDataset,
PoisonLabelDataset,
RangeRatioDataset,
RangeRatioPoisonLabelDataset,
RatioDataset,
RatioPoisonLabelDataset,
)
from anti_kd_backdoor.data.dataset.types import XY_TYPE
class FakeDataset(DatasetInterface, Dataset):
cache: dict[tuple, tuple[torch.Tensor, list[int]]] = dict()
def __init__(self,
*,
x_shape: tuple[int, int, int] = (3, 32, 32),
y_range: tuple[int, int] = (0, 9),
nums: int = 10000,
cache_xy: bool = False) -> None:
self._nums = nums
self._x_shape = x_shape
self._y_range = y_range
self._raw_num_classes = y_range[1] - y_range[0] + 1
if cache_xy:
cache_key = (x_shape, y_range, nums)
if cache_key not in FakeDataset.cache:
x, y = self._prepare_xy()
FakeDataset.cache[cache_key] = (x, y)
x, y = FakeDataset.cache[cache_key]
else:
x, y = self._prepare_xy()
self.data, self.targets = x, y
def __len__(self) -> int:
return len(self.targets)
def __getitem__(self, index: int) -> tuple[torch.Tensor, torch.Tensor]:
x = self.data[index]
y = self.targets[index]
return x, y
def get_xy(self) -> XY_TYPE:
return list(self.data), self.targets.copy()
def set_xy(self, xy: XY_TYPE) -> None:
x, y = xy
assert len(x) == len(y)
self.data = np.stack(x, axis=0)
self.targets = y.copy()
@property
def num_classes(self) -> int:
return len(set(self.targets))
@property
def raw_num_classes(self) -> int:
return self._raw_num_classes
def _prepare_xy(self) -> tuple[torch.Tensor, list[int]]:
x = torch.rand((self._nums, *self._x_shape))
num_per_class = self._nums // self._raw_num_classes
y = [
i for _ in range(num_per_class)
for i in range(self._raw_num_classes)
]
y.extend([self._y_range[0]] * (self._nums - len(y)))
return x, y
class IndexFakeDataset(FakeDataset, IndexDataset):
def __init__(self, *, start_idx: int, end_idx: int, **kwargs) -> None:
FakeDataset.__init__(self, **kwargs)
IndexDataset.__init__(self, start_idx=start_idx, end_idx=end_idx)
class RatioFakeDataset(FakeDataset, RatioDataset):
def __init__(self, *, ratio: float, **kwargs) -> None:
FakeDataset.__init__(self, **kwargs)
RatioDataset.__init__(self, ratio=ratio)
class RangeRatioFakeDataset(FakeDataset, RangeRatioDataset):
def __init__(self, *, range_ratio: tuple[float, float], **kwargs) -> None:
FakeDataset.__init__(self, **kwargs)
RangeRatioDataset.__init__(self, range_ratio=range_ratio)
class IndexRatioFakeDataset(FakeDataset, IndexRatioDataset):
def __init__(self, *, start_idx: int, end_idx: int, ratio: float,
**kwargs) -> None:
FakeDataset.__init__(self, **kwargs)
IndexRatioDataset.__init__(self,
start_idx=start_idx,
end_idx=end_idx,
ratio=ratio)
class PoisonLabelFakeDataset(FakeDataset, PoisonLabelDataset):
def __init__(self, *, poison_label: int, **kwargs) -> None:
FakeDataset.__init__(self, **kwargs)
PoisonLabelDataset.__init__(self, poison_label=poison_label)
class RatioPoisonLabelFakeDataset(FakeDataset, RatioPoisonLabelDataset):
def __init__(self, *, ratio: float, poison_label: int, **kwargs) -> None:
FakeDataset.__init__(self, **kwargs)
RatioPoisonLabelDataset.__init__(self,
ratio=ratio,
poison_label=poison_label)
class RangeRatioPoisonLabelFakeDataset(FakeDataset,
RangeRatioPoisonLabelDataset):
def __init__(self, *, range_ratio: tuple[float, float], poison_label: int,
**kwargs) -> None:
FakeDataset.__init__(self, **kwargs)
RangeRatioPoisonLabelDataset.__init__(self,
range_ratio=range_ratio,
poison_label=poison_label)
FAKE_DATASETS_MAPPING = {
'FakeDataset': FakeDataset,
'IndexFakeDataset': IndexFakeDataset,
'RatioFakeDataset': RatioFakeDataset,
'RangeRatioFakeDataset': RangeRatioFakeDataset,
'IndexRatioFakeDataset': IndexRatioFakeDataset,
'PoisonLabelFakeDataset': PoisonLabelFakeDataset,
'RatioPoisonLabelFakeDataset': RatioPoisonLabelFakeDataset,
'RangeRatioPoisonLabelFakeDataset': RangeRatioPoisonLabelFakeDataset
}
def build_fake_dataset(dataset_cfg: dict) -> FakeDataset:
if 'type' not in dataset_cfg:
raise ValueError('Dataset config must have `type` field')
dataset_type = dataset_cfg.pop('type')
if dataset_type not in FAKE_DATASETS_MAPPING:
raise ValueError(
f'Dataset `{dataset_type}` is not support, '
f'available datasets: {list(FAKE_DATASETS_MAPPING.keys())}')
dataset = FAKE_DATASETS_MAPPING[dataset_type]
return dataset(**dataset_cfg)
|