def load_seed_data(self) -> dict[str, torch.Tensor]:
seed_path = os.path.join(self.folder_path, f'seed_{self.seed_num}.npy')
seed_data: dict[str, torch.Tensor] = {}
seed_data = np.load(seed_path, allow_pickle=True).item() if os.path.exists(seed_path) \
else self.save_seed_data()
seed_data['input'] = to_tensor(seed_data['input'])
seed_data['label'] = to_tensor(seed_data['label'], dtype=torch.long)
return seed_data
def load_seed_data(self) -> dict[str, torch.Tensor]:
seed_path = f'{env["result_dir"]}{self.dataset.name}/{self.name}_{self.seed_num}.npy'
seed_data: dict[str, torch.Tensor] = {}
seed_data = np.load(seed_path, allow_pickle=True).item() if os.path.exists(seed_path) \
else self.save_seed_data()
seed_data['input'] = to_tensor(seed_data['input'])
seed_data['label'] = to_tensor(seed_data['label'], dtype=torch.long)
return seed_data
def load(self, path: str = None):
if path is None:
path = self.folder_path + self.get_filename() + '.npz'
_dict = np.load(path)
self.attack.mark.mark = to_tensor(_dict['mark_list'][self.target_class])
self.attack.mark.alpha_mask = to_tensor(_dict['mask_list'][self.target_class])
self.attack.mark.mask = torch.ones_like(self.attack.mark.mark, dtype=torch.bool)
self.attack.mark.random_pos = False
self.attack.mark.height_offset = 0
self.attack.mark.width_offset = 0
def get_data(self, data: tuple[torch.Tensor], v_noise: float = None, mode='train'):
if v_noise is None:
v_noise = self.v_noise
_input = data[0]
if mode == 'train':
# future warning: to_tensor, to_valid_img
noise: torch.Tensor = torch.normal(mean=0.0, std=v_noise, size=_input.shape)
data[0] = (_input + noise).clamp(0.0, 1.0)
data[1] = _input.detach()
else:
data[0] = _input.detach()
data[1] = _input.clone().detach()
return to_tensor(data[0]), to_tensor(data[1])
def load(self, path: str = None):
if path is None:
path = self.folder_path + self.get_filename(
target_class=self.target_class) + '_best.npy'
_dict = np.load(path, allow_pickle=True).item()
self.attack.mark.mark = to_tensor(_dict[self.target_class]['mark'])
self.attack.mark.alpha_mask = to_tensor(
_dict[self.target_class]['mask'])
self.attack.mark.mask = torch.ones_like(self.attack.mark.mark,
dtype=torch.bool)
self.attack.mark.random_pos = False
self.attack.mark.height_offset = 0
self.attack.mark.width_offset = 0
def load(self, path: str = None):
if path is None:
path = os.path.join(self.folder_path, self.get_filename() + '.npz')
_dict = np.load(path)
self.attack.mark.mark = to_tensor(
_dict['mark_list'][self.target_class])
self.attack.mark.alpha_mask = to_tensor(
_dict['mask_list'][self.target_class])
self.attack.mark.mask = torch.ones_like(self.attack.mark.mark,
dtype=torch.bool)
self.attack.mark.random_pos = False
self.attack.mark.height_offset = 0
self.attack.mark.width_offset = 0
print('defense results loaded from: ', path)
def preprocess_mark(self, data: dict[str, tuple[torch.Tensor,
torch.Tensor]]):
other_x, _ = data['other']
other_set = TensorDataset(other_x)
other_loader = self.dataset.get_dataloader(mode='train',
dataset=other_set,
num_workers=0)
atanh_mark = torch.randn_like(self.mark.mark) * self.mark.mask
atanh_mark.requires_grad_()
self.mark.mark = tanh_func(atanh_mark)
optimizer = optim.Adam([atanh_mark], lr=self.preprocess_lr)
optimizer.zero_grad()
losses = AverageMeter('Loss', ':.4e')
for _epoch in range(self.preprocess_epoch):
loader = other_loader
for (batch_x, ) in loader:
poison_x = self.mark.add_mark(to_tensor(batch_x))
loss = self.loss_mse(poison_x)
loss.backward()
optimizer.step()
optimizer.zero_grad()
self.mark.mark = tanh_func(atanh_mark)
losses.update(loss.item(), n=len(batch_x))
atanh_mark.requires_grad = False
self.mark.mark.detach_()
def __init__(self,
name: str = None,
model: Union[type[_Model], _Model] = _Model,
dataset: Dataset = None,
num_classes: int = None,
folder_path: str = None,
official: bool = False,
pretrain: bool = False,
randomized_smooth: bool = False,
rs_sigma: float = 0.01,
rs_n: int = 100,
suffix: str = '',
**kwargs):
self.param_list: dict[str, list[str]] = {}
self.param_list['model'] = ['folder_path']
if suffix:
self.param_list['model'].append('suffix')
if randomized_smooth:
self.param_list['model'].extend(
['randomized_smooth', 'rs_sigma', 'rs_n'])
self.name: str = name
self.dataset = dataset
self.suffix = suffix
self.randomized_smooth: bool = randomized_smooth
self.rs_sigma: float = rs_sigma
self.rs_n: int = rs_n
self.folder_path = folder_path
if folder_path is not None:
self.folder_path = os.path.normpath(folder_path)
if not os.path.exists(folder_path):
os.makedirs(folder_path)
# ------------Auto-------------- #
loss_weights: np.ndarray = None if 'loss_weights' not in kwargs.keys(
) else kwargs['loss_weights']
if dataset:
if not isinstance(dataset, Dataset):
raise TypeError(f'{type(dataset)=} {dataset=}')
num_classes = num_classes if num_classes is not None else dataset.num_classes
loss_weights = loss_weights if 'loss_weights' in kwargs.keys(
) else dataset.loss_weights
self.num_classes = num_classes # number of classes
self.loss_weights = loss_weights # TODO: what device shall we save loss_weights? numpy, torch, or torch.cuda.
# ------------------------------ #
self.criterion = self.define_criterion(weight=to_tensor(loss_weights))
if isinstance(model, type):
self._model = model(num_classes=num_classes, **kwargs)
else:
self._model = model
self.model = self.get_parallel_model(self._model)
self.activate_params([])
if official:
self.load('official')
if pretrain:
self.load()
self.eval()
if env['num_gpus']:
self.cuda()
def load(self, path: str = None):
if path is None:
path = os.path.join(
self.folder_path,
self.get_filename(target_class=self.target_class) +
'_best.npy')
_dict: dict[str,
dict[str,
torch.Tensor]] = np.load(path,
allow_pickle=True).item()
self.attack.mark.mark = to_tensor(_dict[self.target_class]['mark'])
self.attack.mark.alpha_mask = to_tensor(
_dict[self.target_class]['mask'])
self.attack.mark.mask = torch.ones_like(self.attack.mark.mark,
dtype=torch.bool)
self.attack.mark.random_pos = False
self.attack.mark.height_offset = 0
self.attack.mark.width_offset = 0
print('defense results loaded from: ', path)
def load(self, path: str = None):
if path is None:
path = self.folder_path + self.get_filename() + '.npz'
_dict = np.load(path, allow_pickle=True)
self.attack.mark.mark = to_tensor(_dict['mark_list'][self.attack.target_class])
self.attack.mark.random_pos = False
self.attack.mark.height_offset = 0
self.attack.mark.width_offset = 0
def add_mark_fn(_input, **kwargs):
return _input + self.attack.mark.mark.to(_input.device)
self.attack.mark.add_mark_fn = add_mark_fn
def get_data(self, data: tuple[torch.Tensor, torch.Tensor], org: bool = False, keep_org: bool = True, poison_label=True, **kwargs) -> tuple[torch.Tensor, torch.Tensor]:
if org:
_input, _label = self.model.get_data(data)
else:
_input, _label = self.attack.get_data(data=data, keep_org=keep_org, poison_label=poison_label, **kwargs)
h, w = _input.shape[-2], _input.shape[-1]
_input_list = []
for single_input in _input:
image = to_pil_image(single_input)
image = F.resize(image, (int(h * self.resize_ratio), int(w * self.resize_ratio)), Image.ANTIALIAS)
image = F.resize(image, (h, w))
_input_list.append(to_tensor(image))
return torch.stack(_input_list), _label
def get_loss_weights(self, file_path: str = None, verbose: bool = None) -> np.ndarray:
file_path = file_path if file_path is not None else os.path.join(self.folder_path, 'loss_weights.npy')
if os.path.exists(file_path):
loss_weights = to_tensor(np.load(file_path), dtype='float')
return loss_weights
else:
if verbose:
print('Calculating Loss Weights')
dataset = self.get_full_dataset('train', transform=None)
_, targets = dataset_to_list(dataset, label_only=True)
loss_weights = np.bincount(targets) # TODO: linting problem
assert len(loss_weights) == self.num_classes
loss_weights: np.ndarray = loss_weights.sum() / loss_weights # TODO: linting problem
np.save(file_path, loss_weights)
print('Loss Weights Saved at ', file_path)
return loss_weights
def preprocess_mark(self, data: dict[str, tuple[torch.Tensor,
torch.Tensor]]):
other_x, _ = data['other']
other_set = TensorDataset(other_x)
other_loader = self.dataset.get_dataloader(mode='train',
dataset=other_set,
num_workers=0)
atanh_mark = torch.randn_like(self.mark.mark) * self.mark.mask
atanh_mark.requires_grad_()
self.mark.mark = tanh_func(atanh_mark)
optimizer = optim.Adam([atanh_mark], lr=self.preprocess_lr)
optimizer.zero_grad()
losses = AverageMeter('Loss', ':.4e')
for _epoch in range(self.preprocess_epoch):
# epoch_start = time.perf_counter()
loader = other_loader
# if env['tqdm']:
# loader = tqdm(loader)
for (batch_x, ) in loader:
poison_x = self.mark.add_mark(to_tensor(batch_x))
loss = self.loss_mse(poison_x)
loss.backward()
optimizer.step()
optimizer.zero_grad()
self.mark.mark = tanh_func(atanh_mark)
losses.update(loss.item(), n=len(batch_x))
# epoch_time = str(datetime.timedelta(seconds=int(
# time.perf_counter() - epoch_start)))
# pre_str = '{blue_light}Epoch: {0}{reset}'.format(
# output_iter(_epoch + 1, self.preprocess_epoch), **ansi).ljust(64 if env['color'] else 35)
# _str = ' '.join([
# f'Loss: {losses.avg:.4f},'.ljust(20),
# f'Time: {epoch_time},'.ljust(20),
# ])
# prints(pre_str, _str, prefix='{upline}{clear_line}'.format(**ansi) if env['tqdm'] else '', indent=4)
atanh_mark.requires_grad = False
self.mark.mark.detach_()
def trigger_detect(self, _input: torch.Tensor):
"""
Args:
_input (torch.Tensor): (N, C, H, W)
"""
# get dominant color
dom_c_list = []
for img in _input:
dom_c: torch.Tensor = self.get_dominant_colour(img) # (C)
dom_c_list.append(dom_c)
dom_c = torch.stack(dom_c_list).unsqueeze(-1).unsqueeze(
-1) # (N, C, 1, 1)
# generate random numbers
height, width = _input.shape[-2:]
pos_height: torch.Tensor = torch.randint(
low=0,
high=height - self.size[0],
size=[len(_input), self.sample_num]) # (N, sample_num)
pos_width: torch.Tensor = torch.randint(
low=0,
high=width - self.size[1],
size=[len(_input), self.sample_num]) # (N, sample_num)
pos_list: torch.Tensor = torch.stack(
[pos_height, pos_width]).transpose(0, -1) # (N, sample_num, 2)
# block potential triggers on _input
block_input = _input.unsqueeze(1).repeat(1, self.sample_num, 1, 1,
1) # (N, sample_num, C, H, W)
for i in range(len(_input)):
for j in range(self.sample_num):
x = pos_list[i][j][0]
y = pos_list[i][j][1]
block_input[i, j, :, x:x + self.size[0],
y:y + self.size[1]] = dom_c[i]
# get potential triggers
_input = to_tensor(_input)
block_input = to_tensor(block_input)
org_class = self.model.get_class(_input).unsqueeze(1).repeat(
1, self.sample_num) # (N, sample_num)
block_class_list = []
for i in range(self.sample_num):
block_class = self.model.get_class(
block_input[:, i]) # (N, sample_num)
block_class_list.append(block_class)
block_class = torch.stack(block_class_list, dim=1)
potential_idx: torch.Tensor = org_class.eq(
block_class).detach().cpu() # (N, sample_num)
# confirm triggers
result_list = torch.zeros(len(_input), dtype=torch.bool)
mask_shape = [_input.shape[0], _input.shape[-2], _input.shape[-1]]
mask_list = torch.zeros(mask_shape,
dtype=torch.float) # (N, C, height, width)
mark_class = self.attack.mark
for i in range(len(_input)):
print(f'input {i:3d}')
pos_pairs = pos_list[i][~potential_idx[i]] # (*, 2)
if len(pos_pairs) == 0:
continue
for j, pos in enumerate(pos_pairs):
self.attack.mark.height_offset = pos[0]
self.attack.mark.width_offset = pos[1]
mark_class.org_mark = _input[i, :,
pos[0]:pos[0] + self.size[0],
pos[1]:pos[1] + self.size[1]]
mark_class.org_mask = torch.ones(self.size, dtype=torch.bool)
mark_class.org_alpha_mask = torch.ones(self.size,
dtype=torch.float)
mark_class.mark, mark_class.mask, mark_class.alpha_mask = mark_class.mask_mark(
height_offset=pos[0], width_offset=pos[1])
target_acc = self.confirm_backdoor()
output_str = f' {j:3d} Acc: {target_acc:5.2f}'
if not self.attack.mark.random_pos:
overlap = jaccard_idx(mark_class.mask.detach().cpu(),
self.real_mask.detach().cpu(),
select_num=self.size[0] *
self.size[1])
output_str += f' Jaccard Idx: {overlap:5.3f}'
print(output_str)
if target_acc > self.threshold_t:
result_list[i] = True
mask_list[i] = mark_class.mask
return result_list, mask_list
def nc_filter_img(self) -> torch.Tensor:
h, w = self.dataset.n_dim
mask = torch.ones(h, w, dtype=torch.float)
return to_tensor(mask, non_blocking=False)
def filter_img(self):
h, w = self.dataset.n_dim
mask = torch.zeros(h, w, dtype=torch.float)
mask[2:7, 2:7] = 1
return to_tensor(mask, non_blocking=False)
def save_img(self, img_path: str):
img = self.org_mark * self.org_mask if self.random_pos else self.mark * self.mask
save_tensor_as_img(img_path, img)
def load_npz(self, npz_path: str):
if not os.path.exists(npz_path) and not os.path.exists(
npz_path := os.path.join(dir_path, npz_path)):
raise FileNotFoundError(npz_path.removeprefix(dir_path))
_dict = np.load(npz_path)
if not self.mark_distributed:
self.org_mark = torch.as_tensor(_dict['org_mark'])
self.org_mask = torch.as_tensor(_dict['org_mask'])
self.org_alpha_mask = torch.as_tensor(_dict['org_alpha_mask'])
if not self.random_pos:
self.mark = to_tensor(_dict['mark'])
self.mask = to_tensor(_dict['mask'])
self.alpha_mask = to_tensor(_dict['alpha_mask'])
def save_npz(self, npz_path: str):
_dict = {}
if not self.mark_distributed:
_dict |= {
'org_mark': to_numpy(self.org_mark),
'org_mask': to_numpy(self.org_mask),
'org_alpha_mask': to_numpy(self.org_alpha_mask)
}
if not self.random_pos:
_dict |= {
'mark': to_numpy(self.mark),
'mask': to_numpy(self.mask),
def get_data(data: tuple[torch.Tensor, torch.Tensor],
**kwargs) -> tuple[torch.Tensor, torch.Tensor]:
return to_tensor(data[0]), to_tensor(data[1], dtype='long')
