| 386 | |
| 387 | |
| 388 | class BackdoorDataset: |
| 389 | def __init__( |
| 390 | self, dataset, transform=None, |
| 391 | backdoor_indices=None, trigger_trans=None, attack_trans=None): |
| 392 | |
| 393 | self.dataset = dataset |
| 394 | self.transform = transform |
| 395 | |
| 396 | self.backdoor_indices = set() if ( |
| 397 | backdoor_indices is None) else set(backdoor_indices) |
| 398 | self.trigger_trans = trigger_trans |
| 399 | self.attack_trans = attack_trans |
| 400 | |
| 401 | def __getitem__(self, i): |
| 402 | x, y = self.dataset[i] |
| 403 | |
| 404 | if i in self.backdoor_indices: |
| 405 | x, y = self.trigger_trans(x), self.attack_trans(y) |
| 406 | |
| 407 | if self.transform is not None: |
| 408 | x = self.transform(x) |
| 409 | |
| 410 | return x, y |
| 411 | |
| 412 | def __len__(self): |
| 413 | return len(self.dataset) |
| 414 | |
| 415 | |
| 416 | class IndicesDataset: |