def train_mask_generator(self, verbose: bool = True):
r"""Train :attr:`self.mask_generator`."""
optimizer = torch.optim.Adam(self.mask_generator.parameters(), lr=1e-2, betas=(0.5, 0.9))
lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
optimizer, T_max=self.train_mask_epochs)
loader = self.dataset.loader['train']
dataset = loader.dataset
logger = MetricLogger()
logger.create_meters(loss=None, div=None, norm=None)
print_prefix = 'Mask Epoch'
for _epoch in range(self.train_mask_epochs):
_epoch += 1
idx = torch.randperm(len(dataset))
pos = 0
logger.reset()
header: str = '{blue_light}{0}: {1}{reset}'.format(
print_prefix, output_iter(_epoch, self.train_mask_epochs), **ansi)
header = header.ljust(max(len(print_prefix), 30) + get_ansi_len(header))
self.mask_generator.train()
for data in logger.log_every(loader, header=header) if verbose else loader:
optimizer.zero_grad()
_input, _label = self.model.get_data(data)
batch_size = len(_input)
data2 = sample_batch(dataset, idx=idx[pos:pos + batch_size])
_input2, _label2 = self.model.get_data(data2)
pos += batch_size
_mask = self.get_mask(_input)
_mask2 = self.get_mask(_input2)
input_dist: torch.Tensor = (_input - _input2).flatten(1).norm(p=2, dim=1)
mask_dist: torch.Tensor = (_mask - _mask2).flatten(1).norm(p=2, dim=1) + 1e-5
loss_div = input_dist.div(mask_dist).mean()
loss_norm = _mask.sub(self.mask_density).relu().mean()
loss = self.lambda_norm * loss_norm + self.lambda_div * loss_div
loss.backward()
optimizer.step()
logger.update(n=batch_size, loss=loss.item(), div=loss_div.item(), norm=loss_norm.item())
lr_scheduler.step()
self.mask_generator.eval()
if verbose and (_epoch % (max(self.train_mask_epochs // 5, 1)) == 0 or _epoch == self.train_mask_epochs):
self.validate_mask_generator()
optimizer.zero_grad()
def attack(self, epochs: int, save: bool = False, **kwargs):
# train generator
resnet_model = trojanvision.models.create('resnet18_comp',
dataset=self.dataset,
pretrained=True)
model_extractor = nn.Sequential(*list(resnet_model._model.features.children())[:4])
match self.generator_mode:
case 'resnet':
resnet_model = torchvision.models.resnet18(pretrained=True).to(device=env['device'])
model_extractor = nn.Sequential(*list(resnet_model.children())[:5])
case _:
resnet_model = trojanvision.models.create('resnet18_comp',
dataset=self.dataset,
pretrained=True)
model_extractor = nn.Sequential(*list(resnet_model._model.features.children())[:4])
model_extractor.requires_grad_(False)
model_extractor.train()
self.generator.train()
if self.generator_mode == 'default':
self.generator.requires_grad_()
parameters = self.generator.parameters()
else:
self.generator.bottleneck.requires_grad_()
self.generator.decoder.requires_grad_()
parameters = self.generator[1:].parameters()
optimizer = torch.optim.Adam(parameters, lr=1e-3)
lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
optimizer, T_max=self.train_generator_epochs,
eta_min=1e-5)
logger = MetricLogger()
logger.create_meters(loss=None, acc=None)
for _epoch in range(self.train_generator_epochs):
_epoch += 1
logger.reset()
header: str = '{blue_light}{0}: {1}{reset}'.format(
'Epoch', output_iter(_epoch, self.train_generator_epochs), **ansi)
header = header.ljust(max(len('Epoch'), 30) + get_ansi_len(header))
loader = logger.log_every(self.dataset.loader['train'], header=header,
tqdm_header='Batch')
for data in loader:
optimizer.zero_grad()
_input, _label = self.model.get_data(data)
adv_input = (self.generator(_input) + 1) / 2
_feats = model_extractor(_input)
adv_feats = model_extractor(adv_input)
loss = F.l1_loss(_feats, self.noise_coeff * adv_feats)
loss.backward()
optimizer.step()
batch_size = len(_label)
with torch.no_grad():
org_class = self.model.get_class(_input)
adv_class = self.model.get_class(adv_input)
acc = (org_class == adv_class).float().sum().item() * 100.0 / batch_size
logger.update(n=batch_size, loss=loss.item(), acc=acc)
lr_scheduler.step()
self.save_generator()
optimizer.zero_grad()
self.generator.eval()
self.generator.requires_grad_(False)
self.mark.mark[:-1] = (self.generator(self.mark.mark[:-1].unsqueeze(0))[0] + 1) / 2
self.poison_set = self.get_poison_dataset(load_mark=False)
return super().attack(epochs, save=save, **kwargs)
def optimize_mark(self,
label: int,
loader: Iterable = None,
logger_header: str = '',
verbose: bool = True,
**kwargs) -> tuple[torch.Tensor, float]:
r"""
Args:
label (int): The class label to optimize.
loader (collections.abc.Iterable):
Data loader to optimize trigger.
Defaults to ``self.dataset.loader['train']``.
logger_header (str): Header string of logger.
Defaults to ``''``.
verbose (bool): Whether to use logger for output.
Defaults to ``True``.
**kwargs: Keyword arguments passed to :meth:`loss()`.
Returns:
(torch.Tensor, torch.Tensor):
Optimized mark tensor with shape ``(C + 1, H, W)``
and loss tensor.
"""
atanh_mark = torch.randn_like(self.attack.mark.mark,
requires_grad=True)
optimizer = optim.Adam([atanh_mark],
lr=self.defense_remask_lr,
betas=(0.5, 0.9))
lr_scheduler = optim.lr_scheduler.CosineAnnealingLR(
optimizer, T_max=self.defense_remask_epoch)
optimizer.zero_grad()
loader = loader or self.dataset.loader['train']
# best optimization results
norm_best: float = float('inf')
mark_best: torch.Tensor = None
loss_best: float = None
logger = MetricLogger(indent=4)
logger.create_meters(
loss='{last_value:.3f}',
acc='{last_value:.3f}',
norm='{last_value:.3f}',
entropy='{last_value:.3f}',
)
batch_logger = MetricLogger()
logger.create_meters(loss=None, acc=None, entropy=None)
iterator = range(self.defense_remask_epoch)
if verbose:
iterator = logger.log_every(iterator, header=logger_header)
for _ in iterator:
batch_logger.reset()
for data in loader:
self.attack.mark.mark = tanh_func(atanh_mark) # (c+1, h, w)
_input, _label = self.model.get_data(data)
trigger_input = self.attack.add_mark(_input)
trigger_label = label * torch.ones_like(_label)
trigger_output = self.model(trigger_input)
batch_acc = trigger_label.eq(
trigger_output.argmax(1)).float().mean()
batch_entropy = self.loss(_input,
_label,
target=label,
trigger_output=trigger_output,
**kwargs)
batch_norm: torch.Tensor = self.attack.mark.mark[-1].norm(p=1)
batch_loss = batch_entropy + self.cost * batch_norm
batch_loss.backward()
optimizer.step()
optimizer.zero_grad()
batch_size = _label.size(0)
batch_logger.update(n=batch_size,
loss=batch_loss.item(),
acc=batch_acc.item(),
entropy=batch_entropy.item())
lr_scheduler.step()
self.attack.mark.mark = tanh_func(atanh_mark) # (c+1, h, w)
# check to save best mask or not
loss = batch_logger.meters['loss'].global_avg
acc = batch_logger.meters['acc'].global_avg
norm = float(self.attack.mark.mark[-1].norm(p=1))
entropy = batch_logger.meters['entropy'].global_avg
if norm < norm_best:
mark_best = self.attack.mark.mark.detach().clone()
loss_best = loss
logger.update(loss=loss, acc=acc, norm=norm, entropy=entropy)
if self.check_early_stop(loss=loss,
acc=acc,
norm=norm,
entropy=entropy):
print('early stop')
break
atanh_mark.requires_grad_(False)
self.attack.mark.mark = mark_best
return mark_best, loss_best
def attack(self, epochs: int, optimizer: torch.optim.Optimizer,
lr_scheduler: torch.optim.lr_scheduler._LRScheduler = None,
validate_interval: int = 1, save: bool = False,
verbose: bool = True, **kwargs):
if verbose:
print('train mask generator')
self.mark_generator.requires_grad_(False)
self.mask_generator.requires_grad_()
self.model.requires_grad_(False)
self.train_mask_generator(verbose=verbose)
if verbose:
print()
print('train mark generator and model')
self.mark_generator.requires_grad_()
self.mask_generator.requires_grad_(False)
if not self.natural:
params: list[nn.Parameter] = []
for param_group in optimizer.param_groups:
params.extend(param_group['params'])
self.model.activate_params(params)
mark_optimizer = torch.optim.Adam(self.mark_generator.parameters(), lr=1e-2, betas=(0.5, 0.9))
mark_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
mark_optimizer, T_max=epochs)
loader = self.dataset.loader['train']
dataset = loader.dataset
logger = MetricLogger()
logger.create_meters(loss=None, div=None, ce=None)
if validate_interval != 0:
best_validate_result = self.validate_fn(verbose=verbose)
best_asr = best_validate_result[0]
for _epoch in range(epochs):
_epoch += 1
idx = torch.randperm(len(dataset))
pos = 0
logger.reset()
if not self.natural:
self.model.train()
self.mark_generator.train()
header: str = '{blue_light}{0}: {1}{reset}'.format(
'Epoch', output_iter(_epoch, epochs), **ansi)
header = header.ljust(max(len('Epoch'), 30) + get_ansi_len(header))
for data in logger.log_every(loader, header=header) if verbose else loader:
if not self.natural:
optimizer.zero_grad()
mark_optimizer.zero_grad()
_input, _label = self.model.get_data(data)
batch_size = len(_input)
data2 = sample_batch(dataset, idx=idx[pos:pos + batch_size])
_input2, _label2 = self.model.get_data(data2)
pos += batch_size
final_input, final_label = _input.clone(), _label.clone()
# generate trigger input
trigger_dec, trigger_int = math.modf(len(_label) * self.poison_percent)
trigger_int = int(trigger_int)
if random.uniform(0, 1) < trigger_dec:
trigger_int += 1
x = _input[:trigger_int]
trigger_mark, trigger_mask = self.get_mark(x), self.get_mask(x)
trigger_input = x + trigger_mask * (trigger_mark - x)
final_input[:trigger_int] = trigger_input
final_label[:trigger_int] = self.target_class
# generate cross input
cross_dec, cross_int = math.modf(len(_label) * self.cross_percent)
cross_int = int(cross_int)
if random.uniform(0, 1) < cross_dec:
cross_int += 1
x = _input[trigger_int:trigger_int + cross_int]
x2 = _input2[trigger_int:trigger_int + cross_int]
cross_mark, cross_mask = self.get_mark(x2), self.get_mask(x2)
cross_input = x + cross_mask * (cross_mark - x)
final_input[trigger_int:trigger_int + cross_int] = cross_input
# div loss
if len(trigger_input) <= len(cross_input):
length = len(trigger_input)
cross_input = cross_input[:length]
cross_mark = cross_mark[:length]
cross_mask = cross_mask[:length]
else:
length = len(cross_input)
trigger_input = trigger_input[:length]
trigger_mark = trigger_mark[:length]
trigger_mask = trigger_mask[:length]
input_dist: torch.Tensor = (trigger_input - cross_input).flatten(1).norm(p=2, dim=1)
mark_dist: torch.Tensor = (trigger_mark - cross_mark).flatten(1).norm(p=2, dim=1) + 1e-5
loss_ce = self.model.loss(final_input, final_label)
loss_div = input_dist.div(mark_dist).mean()
loss = loss_ce + self.lambda_div * loss_div
loss.backward()
if not self.natural:
optimizer.step()
mark_optimizer.step()
logger.update(n=batch_size, loss=loss.item(), div=loss_div.item(), ce=loss_ce.item())
if not self.natural and lr_scheduler:
lr_scheduler.step()
mark_scheduler.step()
if not self.natural:
self.model.eval()
self.mark_generator.eval()
if validate_interval != 0 and (_epoch % validate_interval == 0 or _epoch == epochs):
validate_result = self.validate_fn(verbose=verbose)
cur_asr = validate_result[0]
if cur_asr >= best_asr:
best_validate_result = validate_result
best_asr = cur_asr
if save:
self.save()
if not self.natural:
optimizer.zero_grad()
mark_optimizer.zero_grad()
self.mark_generator.requires_grad_(False)
self.mask_generator.requires_grad_(False)
self.model.requires_grad_(False)
return best_validate_result
def train(module: nn.Module,
num_classes: int,
epochs: int,
optimizer: Optimizer,
lr_scheduler: _LRScheduler = None,
lr_warmup_epochs: int = 0,
model_ema: ExponentialMovingAverage = None,
model_ema_steps: int = 32,
grad_clip: float = None,
pre_conditioner: None | KFAC | EKFAC = None,
print_prefix: str = 'Train',
start_epoch: int = 0,
resume: int = 0,
validate_interval: int = 10,
save: bool = False,
amp: bool = False,
loader_train: torch.utils.data.DataLoader = None,
loader_valid: torch.utils.data.DataLoader = None,
epoch_fn: Callable[..., None] = None,
get_data_fn: Callable[..., tuple[torch.Tensor, torch.Tensor]] = None,
forward_fn: Callable[..., torch.Tensor] = None,
loss_fn: Callable[..., torch.Tensor] = None,
after_loss_fn: Callable[..., None] = None,
validate_fn: Callable[..., tuple[float, float]] = None,
save_fn: Callable[..., None] = None,
file_path: str = None,
folder_path: str = None,
suffix: str = None,
writer=None,
main_tag: str = 'train',
tag: str = '',
accuracy_fn: Callable[..., list[float]] = None,
verbose: bool = True,
output_freq: str = 'iter',
indent: int = 0,
change_train_eval: bool = True,
lr_scheduler_freq: str = 'epoch',
backward_and_step: bool = True,
**kwargs):
r"""Train the model"""
if epochs <= 0:
return
get_data_fn = get_data_fn or (lambda x: x)
forward_fn = forward_fn or module.__call__
loss_fn = loss_fn or (lambda _input, _label, _output=None: F.cross_entropy(
_output or forward_fn(_input), _label))
validate_fn = validate_fn or validate
accuracy_fn = accuracy_fn or accuracy
scaler: torch.cuda.amp.GradScaler = None
if not env['num_gpus']:
amp = False
if amp:
scaler = torch.cuda.amp.GradScaler()
best_validate_result = (0.0, float('inf'))
if validate_interval != 0:
best_validate_result = validate_fn(loader=loader_valid,
get_data_fn=get_data_fn,
forward_fn=forward_fn,
loss_fn=loss_fn,
writer=None,
tag=tag,
_epoch=start_epoch,
verbose=verbose,
indent=indent,
**kwargs)
best_acc = best_validate_result[0]
params: list[nn.Parameter] = []
for param_group in optimizer.param_groups:
params.extend(param_group['params'])
len_loader_train = len(loader_train)
total_iter = (epochs - resume) * len_loader_train
logger = MetricLogger()
logger.create_meters(loss=None, top1=None, top5=None)
if resume and lr_scheduler:
for _ in range(resume):
lr_scheduler.step()
iterator = range(resume, epochs)
if verbose and output_freq == 'epoch':
header: str = '{blue_light}{0}: {reset}'.format(print_prefix, **ansi)
header = header.ljust(max(len(header), 30) + get_ansi_len(header))
iterator = logger.log_every(range(resume, epochs),
header=print_prefix,
tqdm_header='Epoch',
indent=indent)
for _epoch in iterator:
_epoch += 1
logger.reset()
if callable(epoch_fn):
activate_params(module, [])
epoch_fn(optimizer=optimizer,
lr_scheduler=lr_scheduler,
_epoch=_epoch,
epochs=epochs,
start_epoch=start_epoch)
loader_epoch = loader_train
if verbose and output_freq == 'iter':
header: str = '{blue_light}{0}: {1}{reset}'.format(
'Epoch', output_iter(_epoch, epochs), **ansi)
header = header.ljust(max(len('Epoch'), 30) + get_ansi_len(header))
loader_epoch = logger.log_every(loader_train,
header=header,
tqdm_header='Batch',
indent=indent)
if change_train_eval:
module.train()
activate_params(module, params)
for i, data in enumerate(loader_epoch):
_iter = _epoch * len_loader_train + i
# data_time.update(time.perf_counter() - end)
_input, _label = get_data_fn(data, mode='train')
if pre_conditioner is not None and not amp:
pre_conditioner.track.enable()
_output = forward_fn(_input, amp=amp, parallel=True)
loss = loss_fn(_input, _label, _output=_output, amp=amp)
if backward_and_step:
optimizer.zero_grad()
if amp:
scaler.scale(loss).backward()
if callable(after_loss_fn) or grad_clip is not None:
scaler.unscale_(optimizer)
if callable(after_loss_fn):
after_loss_fn(_input=_input,
_label=_label,
_output=_output,
loss=loss,
optimizer=optimizer,
loss_fn=loss_fn,
amp=amp,
scaler=scaler,
_iter=_iter,
total_iter=total_iter)
if grad_clip is not None:
nn.utils.clip_grad_norm_(params, grad_clip)
scaler.step(optimizer)
scaler.update()
else:
loss.backward()
if callable(after_loss_fn):
after_loss_fn(_input=_input,
_label=_label,
_output=_output,
loss=loss,
optimizer=optimizer,
loss_fn=loss_fn,
amp=amp,
scaler=scaler,
_iter=_iter,
total_iter=total_iter)
# start_epoch=start_epoch, _epoch=_epoch, epochs=epochs)
if pre_conditioner is not None:
pre_conditioner.track.disable()
pre_conditioner.step()
if grad_clip is not None:
nn.utils.clip_grad_norm_(params, grad_clip)
optimizer.step()
if model_ema and i % model_ema_steps == 0:
model_ema.update_parameters(module)
if _epoch <= lr_warmup_epochs:
# Reset ema buffer to keep copying weights
# during warmup period
model_ema.n_averaged.fill_(0)
if lr_scheduler and lr_scheduler_freq == 'iter':
lr_scheduler.step()
acc1, acc5 = accuracy_fn(_output,
_label,
num_classes=num_classes,
topk=(1, 5))
batch_size = int(_label.size(0))
logger.update(n=batch_size, loss=float(loss), top1=acc1, top5=acc5)
empty_cache()
optimizer.zero_grad()
if lr_scheduler and lr_scheduler_freq == 'epoch':
lr_scheduler.step()
if change_train_eval:
module.eval()
activate_params(module, [])
loss, acc = (logger.meters['loss'].global_avg,
logger.meters['top1'].global_avg)
if writer is not None:
from torch.utils.tensorboard import SummaryWriter
assert isinstance(writer, SummaryWriter)
writer.add_scalars(main_tag='Loss/' + main_tag,
tag_scalar_dict={tag: loss},
global_step=_epoch + start_epoch)
writer.add_scalars(main_tag='Acc/' + main_tag,
tag_scalar_dict={tag: acc},
global_step=_epoch + start_epoch)
if validate_interval != 0 and (_epoch % validate_interval == 0
or _epoch == epochs):
validate_result = validate_fn(module=module,
num_classes=num_classes,
loader=loader_valid,
get_data_fn=get_data_fn,
forward_fn=forward_fn,
loss_fn=loss_fn,
writer=writer,
tag=tag,
_epoch=_epoch + start_epoch,
verbose=verbose,
indent=indent,
**kwargs)
cur_acc = validate_result[0]
if cur_acc >= best_acc:
best_validate_result = validate_result
if verbose:
prints('{purple}best result update!{reset}'.format(**ansi),
indent=indent)
prints(
f'Current Acc: {cur_acc:.3f} '
f'Previous Best Acc: {best_acc:.3f}',
indent=indent)
best_acc = cur_acc
if save:
save_fn(file_path=file_path,
folder_path=folder_path,
suffix=suffix,
verbose=verbose)
if verbose:
prints('-' * 50, indent=indent)
module.zero_grad()
return best_validate_result