def mix_dataset(self, poison_label: bool = True) -> torch.utils.data.Dataset:
clean_dataset = self.dataset.loader['train'].dataset
_input, _label = dataset_to_list(clean_dataset)
_input = torch.stack(_input[self.poison_percent * len(clean_dataset)])
if poison_label:
_label = [self.target_class] * len(_label)
poison_input = self.add_mark(_input)
poison_dataset = TensorListDataset(poison_input, _label)
return torch.utils.data.ConcatDataset([clean_dataset, poison_dataset])
def get_class_set(dataset: torch.utils.data.Dataset, classes: list[int]) -> torch.utils.data.Subset:
indices = np.arange(len(dataset))
if isinstance(dataset, torch.utils.data.Subset):
idx = np.array(dataset.indices)
indices = idx[indices]
dataset = dataset.dataset
_, targets = dataset_to_list(dataset=dataset, label_only=True)
idx_bool = np.isin(targets, classes)
idx = np.arange(len(dataset))[idx_bool]
idx = np.intersect1d(idx, indices)
return torch.utils.data.Subset(dataset, idx)
def mix_dataset(self) -> torch.utils.data.Dataset:
clean_dataset = self.dataset.loader['train'].dataset
subset, otherset = ImageSet.split_dataset(clean_dataset, percent=self.poison_percent)
if not len(subset):
return clean_dataset
_input, _label = dataset_to_list(subset)
_input = torch.stack(_input)
_label = torch.tensor(_label, dtype=torch.long)
_label += torch.randint_like(_label, low=1, high=self.model.num_classes)
_label %= self.model.num_classes
_label = _label.tolist()
poison_dataset = TensorListDataset(_input, _label)
return torch.utils.data.ConcatDataset([poison_dataset, otherset])
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 initialize_npz(self,
mode_list: list[str] = ['train', 'valid'],
transform: transforms.Lambda = transforms.Lambda(
lambda x: np.array(x)),
**kwargs):
if not self.valid_set:
mode_list.remove('valid')
json_path = os.path.join(self.folder_path, 'class_to_idx.json')
for mode in mode_list:
npz_path = os.path.join(self.folder_path,
f'{self.name}_{mode}.npz')
if not os.path.exists(npz_path):
print('{yellow}initialize npz{reset}: '.format(**ansi),
npz_path)
dataset: datasets.ImageFolder = self.get_org_dataset(
mode, transform=transform, data_format='folder')
data, targets = dataset_to_list(dataset)
data = np.stack(data)
np.savez(npz_path, data=data, targets=targets)
with open(json_path, 'w') as f:
json.dump(dataset.class_to_idx, f)
print('{green}initialize npz finish{reset}: '.format(**ansi))
import numpy as np
from sklearn import metrics
from trojanzoo.utils.data import dataset_to_list
from art.estimators.classification import PyTorchClassifier # type: ignore
classifier = PyTorchClassifier(
model=model._model,
loss=model.criterion,
input_shape=dataset.data_shape,
nb_classes=model.num_classes,
)
model._validate()
from art.attacks.inference.membership_inference import LabelOnlyDecisionBoundary as Attack # type: ignore
attack = Attack(classifier)
x_train, y_train = dataset_to_list(dataset.get_dataset('train'))
x_train, y_train = to_numpy(torch.stack(x_train)), to_numpy(y_train)
x_valid, y_valid = dataset_to_list(dataset.get_dataset('valid'))
x_valid, y_valid = to_numpy(torch.stack(x_valid)), to_numpy(y_valid)
sample_size = 64
tau_path = os.path.normpath(os.path.join(model.folder_path,
model.name)) + '_tau'
t_idx = np.arange(len(x_train))
v_idx = np.arange(len(x_valid))
np.random.shuffle(t_idx)
np.random.shuffle(v_idx)
dataset = trojanvision.datasets.create(**args.__dict__)
model = trojanvision.models.create(dataset=dataset, **args.__dict__)
if env['verbose']:
summary(env=env, dataset=dataset, model=model)
model._validate()
print('\n\n')
from art.estimators.classification import PyTorchClassifier # type: ignore
classifier = PyTorchClassifier(
model=model._model,
loss=model.criterion,
input_shape=dataset.data_shape,
nb_classes=model.num_classes,
)
x_train, y_train = dataset_to_list(dataset.get_dataset('train'))
x_train, y_train = to_numpy(torch.stack(x_train)), to_numpy(y_train)
# valid_train, valid_valid = dataset.split_set(dataset.get_dataset('valid'), length=5000)
# x_train, y_train = dataset_to_list(valid_train)
# x_train, y_train = to_numpy(torch.stack(x_train)), to_numpy(y_train)
# valid_loader = dataset.get_dataloader('valid', dataset=valid_valid)
# thieved_model._validate(print_prefix='Before Stealing', loader=valid_loader)
# thieved_model._validate(print_prefix='After Stealing', loader=valid_loader)
import art.attacks.extraction # type:ignore
for name in ['CopycatCNN', 'KnockoffNets']:
for use_probability in [True, False]:
print(name + (' Probabilistic' if use_probability else ' Argmax'))
AttackClass = getattr(art.attacks.extraction, name)
if __name__ == '__main__':
parser = argparse.ArgumentParser()
trojanvision.environ.add_argument(parser)
trojanvision.datasets.add_argument(parser)
trojanvision.models.add_argument(parser)
args = parser.parse_args()
env = trojanvision.environ.create(**args.__dict__)
dataset = trojanvision.datasets.create(**args.__dict__)
model = trojanvision.models.create(dataset=dataset, **args.__dict__)
if env['verbose']:
summary(env=env, dataset=dataset, model=model)
model._validate()
train_data, _ = dataset_to_list(dataset.get_dataset(mode='train'))
valid_data, _ = dataset_to_list(dataset.get_dataset(mode='valid'))
train_data, valid_data = torch.stack(train_data), torch.stack(valid_data)
train_data, valid_data = train_data.to(env['device']), valid_data.to(
env['device'])
train_feats, valid_feats = get_feats(train_data,
model), get_feats(valid_data, model)
total_feats = torch.cat([train_feats, valid_feats])
Y_total = total_feats.mean(dim=0)
bmss: float = 0.0
wmss: float = 0.0
for c in range(dataset.num_classes):
print('class: ', c)
train_data, _ = dataset_to_list(
dataset.get_dataset(mode='train', class_list=[c]))