def get_transferset(self): """ :return: """ # for rho_current in range(self.rho): rho_current = 0 while self.blackbox.call_count < self.budget: print('=> Beginning substitute epoch {} (|D| = {})'.format(rho_current, len(self.D))) # -------------------------- 0. Initialize Model model_adv = zoo.get_net(self.model_adv_name, self.modelfamily, self.model_adv_pretrained, num_classes=self.num_classes) model_adv = model_adv.to(self.device) # -------------------------- 1. Train model on D model_adv = model_utils.train_model(model_adv, self.D, self.out_dir, num_workers=10, checkpoint_suffix='.{}'.format(self.blackbox.call_count), device=self.device, epochs=self.train_epochs, log_interval=500, lr=0.1, momentum=0.9, batch_size=self.batch_size, lr_gamma=0.1, testset=self.testset, criterion_train=model_utils.soft_cross_entropy) # -------------------------- 2. Evaluate model # _, acc = model_utils.test_step(model_adv, self.testloader, nn.CrossEntropyLoss(reduction='mean'), # device=self.device, epoch=rho_current) # self.accuracies.append(acc) # -------------------------- 3. Jacobian-based data augmentation if self.aug_strategy in ['jbda', 'jbself']: self.D = self.jacobian_augmentation(model_adv, rho_current) elif self.aug_strategy == 'jbtop{}'.format(self.topk): self.D = self.jacobian_augmentation_topk(model_adv, rho_current) else: raise ValueError('Unrecognized augmentation strategy: "{}"'.format(self.aug_strategy)) # -------------------------- 4. End if necessary rho_current += 1 if (self.blackbox.call_count >= self.budget) or ((self.rho is not None) and (rho_current >= self.rho)): print('=> # BB Queries ({}) >= budget ({}). Ending attack.'.format(self.blackbox.call_count, self.budget)) model_adv = zoo.get_net(self.model_adv_name, self.modelfamily, self.model_adv_pretrained, num_classes=self.num_classes) model_adv = model_adv.to(self.device) model_adv = model_utils.train_model(model_adv, self.D, self.out_dir, num_workers=10, checkpoint_suffix='.{}'.format(self.blackbox.call_count), device=self.device, epochs=self.final_train_epochs, log_interval=500, lr=0.01, momentum=0.9, batch_size=self.batch_size, lr_gamma=0.1, testset=self.testset, criterion_train=model_utils.soft_cross_entropy) break print() return self.D, model_adv
def from_modeldir(cls, model_dir, vocab_size, num_classes, embed_dim, device=None, output_type='probs'): device = torch.device('cuda') if device is None else device # What was the model architecture used by this model? params_path = osp.join(model_dir, 'params.json') with open(params_path) as jf: params = json.load(jf) model_arch = params['model_arch'] # num_classes = params['num_classes'] victim_dataset = params.get('dataset', 'imagenet') modelfamily = datasets.dataset_to_modelfamily[victim_dataset] # Instantiate the model # model = model_utils.get_net(model_arch, n_output_classes=num_classes) model = zoo.get_net(model_arch, modelfamily, pretrained=None, vocab_size=vocab_size, num_class=num_classes, embed_dim=embed_dim) model = model.to(device) # Load weights checkpoint_path = osp.join(model_dir, 'model_best.pth.tar') if not osp.exists(checkpoint_path): checkpoint_path = osp.join(model_dir, 'checkpoint.pth.tar') print("=> loading checkpoint '{}'".format(checkpoint_path)) checkpoint = torch.load(checkpoint_path, map_location=device) #To run on local machine. Remove when you run on the server. epoch = checkpoint['epoch'] best_test_acc = checkpoint['best_acc'] model.load_state_dict(checkpoint['state_dict']) print("=> loaded checkpoint (epoch {}, acc={:.2f})".format(epoch, best_test_acc)) blackbox = cls(model, device, output_type) return blackbox
def from_modeldir(cls, model_dir, device=None, output_type='probs', **kwargs): device = torch.device('cuda') if device is None else device # What was the model architecture used by this model? params_path = osp.join(model_dir, 'params.json') with open(params_path) as jf: params = json.load(jf) model_arch = params['model_arch'] num_classes = params['num_classes'] if 'queryset' in params: dataset_name = params['queryset'] elif 'testdataset' in params: dataset_name = params['testdataset'] elif 'dataset' in params: dataset_name = params['dataset'] else: raise ValueError('Unable to determine model family') modelfamily = datasets.dataset_to_modelfamily[dataset_name] # Instantiate the model model = zoo.get_net(model_arch, modelfamily, num_classes=num_classes) model = model.to(device) # Load weights checkpoint_path = osp.join(model_dir, 'model_best.pth.tar') if not osp.exists(checkpoint_path): checkpoint_path = osp.join(model_dir, 'checkpoint.pth.tar') print("=> loading checkpoint '{}'".format(checkpoint_path)) checkpoint = torch.load(checkpoint_path) epoch = checkpoint['epoch'] best_test_acc = checkpoint['best_acc'] model.load_state_dict(checkpoint['state_dict']) print("=> loaded checkpoint (epoch {}, acc={:.2f})".format( epoch, best_test_acc)) # print(cls, model, device, output_type, kwargs) blackbox = cls(model=model, device=device, output_type=output_type, dataset_name=dataset_name, modelfamily=modelfamily, model_arch=model_arch, num_classes=num_classes, model_dir=model_dir, **kwargs) return blackbox
def __init__(self, epsilon=None, optim='linesearch', model_adv_proxy=None, max_grad_layer=None, ydist='l1', oracle='extreme', disable_jacobian=False, objmax=False, out_path=None, log_prefix='', *args, **kwargs): super().__init__(*args, **kwargs) print('=> MAD ({})'.format( [self.dataset_name, epsilon, optim, ydist, oracle])) self.epsilon = epsilon self.out_path = out_path self.disable_jacobian = bool(disable_jacobian) if self.disable_jacobian: print('') print('!!!WARNING!!! Using G = eye(K)') print('') self.objmax = bool(objmax) # Victim's assumption of adversary's model print('Proxy for F_A = ', model_adv_proxy) if model_adv_proxy is not None: if osp.isdir(model_adv_proxy): model_adv_proxy_params = osp.join(model_adv_proxy, 'params.json') model_adv_proxy = osp.join(model_adv_proxy, 'checkpoint.pth.tar') with open(model_adv_proxy_params, 'r') as rf: proxy_params = json.load(rf) model_adv_proxy_arch = proxy_params['model_arch'] print('Loading proxy ({}) parameters: {}'.format( model_adv_proxy_arch, model_adv_proxy)) assert osp.exists(model_adv_proxy), 'Does not exist: {}'.format( model_adv_proxy) self.model_adv_proxy = zoo.get_net(model_adv_proxy_arch, self.modelfamily, pretrained=model_adv_proxy, num_classes=self.num_classes) self.model_adv_proxy = self.model_adv_proxy.to(self.device) else: self.model_adv_proxy = self.model # To compute stats self.dataset = datasets.__dict__[self.dataset_name] self.modelfamily = datasets.dataset_to_modelfamily[self.dataset_name] self.train_transform = datasets.modelfamily_to_transforms[ self.modelfamily]['train'] self.test_transform = datasets.modelfamily_to_transforms[ self.modelfamily]['test'] self.testset = self.dataset(train=False, transform=self.test_transform) self.K = len(self.testset.classes) self.D = None self.ydist = ydist assert ydist in ['l1', 'l2', 'kl'] # Which oracle to use self.oracle = oracle assert self.oracle in ['extreme', 'random', 'argmin', 'argmax'] # Which algorithm to use to optimize self.optim = optim assert optim in ['linesearch', 'projections', 'greedy'] # Gradients from which layer to use? assert max_grad_layer in [None, 'all'] self.max_grad_layer = max_grad_layer # Track some data for debugging self.queries = [] # List of (x_i, y_i, y_i_prime, distance) self.log_path = osp.join(out_path, 'distance{}.log.tsv'.format(log_prefix)) if not osp.exists(self.log_path): with open(self.log_path, 'w') as wf: columns = [ 'call_count', 'l1_mean', 'l1_std', 'l2_mean', 'l2_std', 'kl_mean', 'kl_std' ] wf.write('\t'.join(columns) + '\n') self.jacobian_times = []
def __init__(self, epsilon=None, optim='linesearch', model_adv_proxy=None, max_grad_layer=None, ydist='l1', oracle='extreme', model_adv=None, model_adv_proxy_notrain=False, out_path=None, disable_jacobian=False, attacker_argmax=False, adv_optimizer='sgd', objmax=False, log_prefix='', *args, **kwargs): super().__init__(*args, **kwargs) print('=> MAD ({})'.format([self.dataset_name, epsilon, optim, ydist, oracle])) self.epsilon = epsilon self.out_path = out_path self.disable_jacobian = bool(disable_jacobian) if self.disable_jacobian: print('') print('!!!WARNING!!! Using G = eye(K)') print('') self.attacker_argmax = bool(attacker_argmax) if self.attacker_argmax: print('') print('!!!WARNING!!! Argmax of perturbed probabilities used to train attacker model') print('') self.objmax = bool(objmax) ''' Here, we refer to three models: a) self.model: Victim's model (already initialized by super) b) self.model_adv: Adversary's model c) self.model_adv_proxy: Proxy to adversary's model ''' '''(b) self.model_adv: Adversary's model ''' # Unlike BreakSGD, here we assume a perfect-knowledge adversary - which will also be trained online # Initialize the adv model print('=> Initializing adv model compatible with: ', self.model_arch, self.modelfamily, self.dataset_name, self.num_classes) if model_adv is not None and osp.isdir(model_adv): model_adv = osp.join(model_adv, 'checkpoint.pth.tar') assert model_adv is None or osp.exists(model_adv) self.model_adv = zoo.get_net(self.model_arch, self.modelfamily, pretrained=model_adv, num_classes=self.num_classes) self.adv_optimizer = adv_optimizer assert adv_optimizer in ['sgd', 'sgdm', 'adam'] if self.adv_optimizer == 'sgd': self.model_adv_optimizer = torch.optim.SGD(self.model_adv.parameters(), lr=0.1 / 64) elif self.adv_optimizer == 'sgdm': self.model_adv_optimizer = torch.optim.SGD(self.model_adv.parameters(), lr=0.1 / 64, momentum=0.5) elif self.adv_optimizer == 'adam': self.model_adv_optimizer = torch.optim.Adam(self.model_adv.parameters(), lr=0.001 / 64) else: raise ValueError('Unrecognized optimizer') self.model_adv = self.model_adv.to(self.device) self.model_adv.train() '''(c) self.model_adv_proxy: Proxy to adversary's model''' self.model_adv_proxy_notrain = bool(model_adv_proxy_notrain) if model_adv_proxy is None: # Perfect Knowledge (White-box attacker) self.model_adv_proxy = self.model_adv else: if osp.isdir(model_adv_proxy): model_adv_proxy = osp.join(model_adv_proxy, 'checkpoint.pth.tar') assert osp.exists(model_adv_proxy), 'Does not exist: {}'.format(model_adv_proxy) print('=== Models used for experiment ===') print('F_V : ', self.model_dir) print('F_A : ', osp.dirname(model_adv)) print('F_A (proxy): ', osp.dirname(model_adv_proxy)) print('F_A (proxy) trained online?: ', not self.model_adv_proxy_notrain) print('==================================') # Gray-box attacker self.model_adv_proxy = zoo.get_net(self.model_arch, self.modelfamily, pretrained=model_adv_proxy, num_classes=self.num_classes) if self.adv_optimizer == 'sgd': self.model_adv_proxy_optimizer = torch.optim.SGD(self.model_adv.parameters(), lr=0.1 / 64) elif self.adv_optimizer == 'sgdm': self.model_adv_proxy_optimizer = torch.optim.SGD(self.model_adv.parameters(), lr=0.1 / 64, momentum=0.5) elif self.adv_optimizer == 'adam': self.model_adv_proxy_optimizer = torch.optim.Adam(self.model_adv.parameters()) else: raise ValueError('Unrecognized optimizer') self.model_adv_proxy = self.model_adv_proxy.to(self.device) # To compute stats self.dataset = datasets.__dict__[self.dataset_name] self.modelfamily = datasets.dataset_to_modelfamily[self.dataset_name] self.train_transform = datasets.modelfamily_to_transforms[self.modelfamily]['train'] self.test_transform = datasets.modelfamily_to_transforms[self.modelfamily]['test'] self.trainset = self.dataset(train=True, transform=self.train_transform) self.testset = self.dataset(train=False, transform=self.test_transform) self.test_loader = DataLoader(self.testset, batch_size=128, shuffle=False, num_workers=5) self.best_test_acc = 0. # Also keep a mini-testset to eval victim model using current strategy self.minitestset = self.dataset(train=False, transform=self.test_transform) self.minitestset = torch.utils.data.Subset(self.minitestset, indices=np.arange(1000)) self.minitest_loader = DataLoader(self.minitestset, batch_size=1, shuffle=False, num_workers=1) self.K = len(self.testset.classes) self.D = None self.ydist = ydist assert ydist in ['l1', 'l2', 'kl'] # Which oracle to use self.oracle = oracle assert self.oracle in ['extreme', 'random', 'argmin', 'argmax'] # Which algorithm to use to optimize self.optim = optim assert optim in ['linesearch', 'projections', 'greedy'] # Gradients from which layer to use? assert max_grad_layer in [None, 'all'] self.max_grad_layer = max_grad_layer # Track some data for debugging self.queries = [] # List of (x_i, y_i, y_i_prime, distance) self.run_id = str(datetime.now()) self.log_path = osp.join(out_path, 'online.log.tsv') if not osp.exists(self.log_path): with open(self.log_path, 'w') as wf: columns = ['run_id', 'epoch', 'split', 'loss', 'accuracy', 'best_accuracy', 'l1_mean', 'l1_std', 'l2_mean', 'l2_std', 'kl_mean', 'kl_std'] wf.write('\t'.join(columns) + '\n')
def main(): parser = argparse.ArgumentParser(description='Train a model') # Required arguments parser.add_argument('dataset', metavar='DS_NAME', type=str, help='Dataset name') parser.add_argument('model_arch', metavar='MODEL_ARCH', type=str, help='Model name') # Optional arguments parser.add_argument('-o', '--out_path', metavar='PATH', type=str, help='Output path for model', default=cfg.MODEL_DIR) parser.add_argument('-d', '--device_id', metavar='D', type=int, help='Device id. -1 for CPU.', default=0) parser.add_argument('-b', '--batch-size', type=int, default=64, metavar='N', help='input batch size for training (default: 64)') parser.add_argument('-e', '--epochs', type=int, default=100, metavar='N', help='number of epochs to train (default: 100)') parser.add_argument('--lr', type=float, default=0.1, metavar='LR', help='learning rate (default: 0.1)') parser.add_argument('--momentum', type=float, default=0.5, metavar='M', help='SGD momentum (default: 0.5)') parser.add_argument( '--log-interval', type=int, default=100, metavar='N', help='how many batches to wait before logging training status') parser.add_argument('--resume', default=None, type=str, metavar='PATH', help='path to latest checkpoint (default: none)') parser.add_argument('--lr-step', type=int, default=30, metavar='N', help='Step sizes for LR') parser.add_argument('--lr-gamma', type=float, default=0.1, metavar='N', help='LR Decay Rate') parser.add_argument('-w', '--num_workers', metavar='N', type=int, help='# Worker threads to load data', default=10) parser.add_argument('--train_subset', type=int, help='Use a subset of train set', default=None) parser.add_argument('--pretrained', type=str, help='Use pretrained network', default=None) parser.add_argument('--weighted-loss', action='store_true', help='Use a weighted loss', default=None) args = parser.parse_args() params = vars(args) # torch.manual_seed(cfg.DEFAULT_SEED) if params['device_id'] >= 0: os.environ["CUDA_VISIBLE_DEVICES"] = str(params['device_id']) device = torch.device('cuda') else: device = torch.device('cpu') # ----------- Set up dataset dataset_name = params['dataset'] valid_datasets = datasets.__dict__.keys() if dataset_name not in valid_datasets: raise ValueError( 'Dataset not found. Valid arguments = {}'.format(valid_datasets)) dataset = datasets.__dict__[dataset_name] modelfamily = datasets.dataset_to_modelfamily[dataset_name] train_transform = datasets.modelfamily_to_transforms[modelfamily]['train'] test_transform = datasets.modelfamily_to_transforms[modelfamily]['test'] trainset = dataset(train=True, transform=train_transform) testset = dataset(train=False, transform=test_transform) num_classes = len(trainset.classes) params['num_classes'] = num_classes if params['train_subset'] is not None: idxs = np.arange(len(trainset)) ntrainsubset = params['train_subset'] idxs = np.random.choice(idxs, size=ntrainsubset, replace=False) trainset = Subset(trainset, idxs) # ----------- Set up model model_name = params['model_arch'] pretrained = params['pretrained'] # model = model_utils.get_net(model_name, n_output_classes=num_classes, pretrained=pretrained) model = zoo.get_net(model_name, modelfamily, pretrained, num_classes=num_classes) model = model.to(device) # ----------- Train out_path = params['out_path'] model_utils.train_model(model, trainset, testset=testset, device=device, **params) # Store arguments params['created_on'] = str(datetime.now()) params_out_path = osp.join(out_path, 'params.json') with open(params_out_path, 'w') as jf: json.dump(params, jf, indent=True)
def __init__(self, beta=1.0, gamma=1.0, out_path=None, model_adv=None, attacker_argmax=False, adv_optimizer='sgd', log_prefix='', *args, **kwargs): super().__init__(*args, **kwargs) print('=> ReverseSigmoid ({})'.format([beta, gamma])) assert beta >= 0. assert gamma >= 0. self.beta = beta self.gamma = gamma self.attacker_argmax = bool(attacker_argmax) if self.attacker_argmax: print('') print( '!!!WARNING!!! Argmax of perturbed probabilities used to train attacker model' ) print('') ''' White-box specific stuff ''' self.out_path = out_path # self.model_adv: Adversary's model # Initialize the adv model print('=> Initializing adv model compatible with: ', self.model_arch, self.modelfamily, self.dataset_name, self.num_classes) if model_adv is not None and osp.isdir(model_adv): model_adv = osp.join(model_adv, 'checkpoint.pth.tar') assert model_adv is None or osp.exists(model_adv) self.model_adv = zoo.get_net(self.model_arch, self.modelfamily, pretrained=model_adv, num_classes=self.num_classes) self.adv_optimizer = adv_optimizer assert adv_optimizer in ['sgd', 'sgdm', 'adam'] if self.adv_optimizer == 'sgd': self.model_adv_optimizer = torch.optim.SGD( self.model_adv.parameters(), lr=0.1 / 64) elif self.adv_optimizer == 'sgdm': self.model_adv_optimizer = torch.optim.SGD( self.model_adv.parameters(), lr=0.1 / 64, momentum=0.5) elif self.adv_optimizer == 'adam': self.model_adv_optimizer = torch.optim.Adam( self.model_adv.parameters(), lr=0.001 / 64) else: raise ValueError('Unrecognized optimizer') self.model_adv = self.model_adv.to(self.device) self.model_adv.train() # To compute stats self.dataset = datasets.__dict__[self.dataset_name] self.modelfamily = datasets.dataset_to_modelfamily[self.dataset_name] self.train_transform = datasets.modelfamily_to_transforms[ self.modelfamily]['train'] self.test_transform = datasets.modelfamily_to_transforms[ self.modelfamily]['test'] self.trainset = self.dataset(train=True, transform=self.train_transform) self.testset = self.dataset(train=False, transform=self.test_transform) self.test_loader = DataLoader(self.testset, batch_size=128, shuffle=False, num_workers=5) self.best_test_acc = 0. # Also keep a mini-testset to eval victim model using current strategy self.minitestset = self.dataset(train=False, transform=self.test_transform) self.minitestset = torch.utils.data.Subset(self.minitestset, indices=np.arange(1000)) self.minitest_loader = DataLoader(self.minitestset, batch_size=1, shuffle=False, num_workers=1) # Track some data for debugging self.queries = [] # List of (x_i, y_i, y_i_prime, distance) self.run_id = str(datetime.now()) self.log_path = osp.join(out_path, 'online.log.tsv') if not osp.exists(self.log_path): with open(self.log_path, 'w') as wf: columns = [ 'run_id', 'epoch', 'split', 'loss', 'accuracy', 'best_accuracy', 'l1_mean', 'l1_std', 'l2_mean', 'l2_std', 'kl_mean', 'kl_std' ] wf.write('\t'.join(columns) + '\n')
def main(): parser = argparse.ArgumentParser(description='Train a model') # Required arguments parser.add_argument('model_dir', metavar='DIR', type=str, help='Directory containing transferset.pickle') parser.add_argument('model_arch', metavar='MODEL_ARCH', type=str, help='Model name') parser.add_argument('testdataset', metavar='DS_NAME', type=str, help='Name of test') parser.add_argument('-o', '--output', type=str, help="Output dir") parser.add_argument('--budget', metavar='N', type=int, help='Size of transfer set to construct', required=True) # Optional arguments parser.add_argument('-b', '--batch-size', type=int, default=1024, metavar='N', help='input batch size for training (default: 64)') parser.add_argument('--resume', default=None, type=str, metavar='PATH', help='path to latest checkpoint (default: none)') parser.add_argument('--device_id', default=0, type=int) parser.add_argument('-w', '--num_workers', metavar='N', type=int, help='# Worker threads to load data', default=10) args = parser.parse_args() params = vars(args) random.seed(cfg.DEFAULT_SEED) np.random.seed(cfg.DEFAULT_SEED) torch.cuda.manual_seed(cfg.DEFAULT_SEED) torch.manual_seed(cfg.DEFAULT_SEED) if params['device_id'] >= 0: os.environ["CUDA_VISIBLE_DEVICES"] = str(params['device_id']) device = torch.device('cuda') else: device = torch.device('cpu') model_dir = params['model_dir'] output_dir = params['output'] if not os.path.exists(output_dir): os.makedirs(output_dir) # ----------- Set up transferset transferset_path = osp.join(model_dir, f"transferset_{params['budget']}.pickle") with open(transferset_path, 'rb') as rf: transferset_samples = pickle.load(rf) num_classes = transferset_samples[0][1].size(0) print('=> found transfer set with {} samples, {} classes'.format( len(transferset_samples), num_classes)) # Must transfer target vector to one-hot vector for resnet10 new_transferset_samples = [] print('=> Using argmax labels (instead of posterior probabilities)') for i in range(len(transferset_samples)): x_i, y_i = transferset_samples[i] argmax_k = y_i.argmax() y_i_1hot = torch.zeros_like(y_i) y_i_1hot[argmax_k] = 1. new_transferset_samples.append((x_i, y_i_1hot)) transferset_samples = new_transferset_samples if params['testdataset'] == 'CIFAR10': num_classes = 10 # ----------- Set up testset dataset_name = params['testdataset'] valid_datasets = datasets.__dict__.keys() modelfamily = datasets.dataset_to_modelfamily[dataset_name] transform = datasets.modelfamily_to_transforms[modelfamily]['test'] if dataset_name not in valid_datasets: raise ValueError( 'Dataset not found. Valid arguments = {}'.format(valid_datasets)) dataset = datasets.__dict__[dataset_name] testset = dataset(train=False, transform=transform) if len(testset.classes) != num_classes: raise ValueError( '# Transfer classes ({}) != # Testset classes ({})'.format( num_classes, len(testset.classes))) # transferset = samples_to_transferset(transferset_samples, budget=params['budget'], transform=transform) transferset_fn = partial( samples_to_transferset, samples=transferset_samples, transform=transform, ) # ----------- Set up model model_name = params['model_arch'] # model = model_utils.get_net(model_name, n_output_classes=num_classes, pretrained=pretrained) model = zoo.get_net(model_name, modelfamily, False, num_classes=num_classes) if params['model_arch'] == "resnet10": model = get_transferred_pytorch_model_resnet10() elif params['model_arch'] == "resnet18": model = get_transferred_pytorch_model_resnet18(testset) else: raise NotImplementedError model = model.cuda() train_fn = partial( train_model, epochs=10, batch_size=256, log_interval=50, lr=5e-3, momentum=0.5, lr_step=30, lr_gamma=0.1, ) # Protect all convs test_prune_allconv_ratio_subset(model=model, train_fn=train_fn, transferset_fn=transferset_fn, testset=testset, out_path=output_dir, **params)
def main(): parser = argparse.ArgumentParser(description='Train a model') # Required arguments parser.add_argument('dataset', metavar='DS_NAME', type=str, help='Dataset name') parser.add_argument('model_arch', metavar='MODEL_ARCH', type=str, help='Model name') # Optional arguments parser.add_argument('-o', '--out_path', metavar='PATH', type=str, help='Output path for model', default=cfg.MODEL_DIR) parser.add_argument('-d', '--device_id', metavar='D', type=int, help='Device id. -1 for CPU.', default=0) parser.add_argument('-b', '--batch-size', type=int, default=32, metavar='N', help='input batch size for training (default: 64)') parser.add_argument('--datadir', default='.data', help='data directory (default=.data)') parser.add_argument('-e', '--epochs', type=int, default=10, metavar='N', help='number of epochs to train (default: 100)') parser.add_argument('--embed-dim', type=int, default=32, help='embed dim. (default=32)') parser.add_argument('--lr', type=float, default=4.0, metavar='LR', help='learning rate (default: 0.1)') # parser.add_argument('--momentum', type=float, default=0.5, metavar='M', # help='SGD momentum (default: 0.5)') parser.add_argument('--log-interval', type=int, default=100, metavar='N', help='how many batches to wait before logging training status') parser.add_argument('--resume', default=None, type=str, metavar='PATH', help='path to latest checkpoint (default: none)') parser.add_argument('--lr-step', type=int, default=1, metavar='N', help='Step sizes for LR') parser.add_argument('--lr-gamma', type=float, default=0.8, metavar='N', help='LR Decay Rate') parser.add_argument('-w', '--num_workers', metavar='N', type=int, help='# Worker threads to load data', default=10) parser.add_argument('--train_subset', type=int, help='Use a subset of train set', default=None) parser.add_argument('--pretrained', type=str, help='Use pretrained network', default=None) parser.add_argument('--weighted-loss', action='store_true', help='Use a weighted loss', default=None) args = parser.parse_args() params = vars(args) # torch.manual_seed(cfg.DEFAULT_SEED) # extract parameter arguments into variables embed_dim = params['embed_dim'] dataset_name = params['dataset'] datadir = params['datadir'] out_path = params['out_path'] batch_size = params['batch_size'] lr = params['lr'] lr_gamma = params['lr_gamma'] num_workers = params['num_workers'] num_epochs = params['epochs'] model_name = params['model_arch'] pretrained = params['pretrained'] if params['device_id'] >= 0: os.environ["CUDA_VISIBLE_DEVICES"] = str(params['device_id']) device = torch.device('cuda') else: device = torch.device('cpu') # Currently supports only the torchtext datasets valid_datasets = list(text_classification.DATASETS.keys()) if dataset_name not in valid_datasets: raise ValueError('Dataset not found. Valid arguments = {}'.format(valid_datasets)) # Currently, we will support classification tasks only. Sentiment task is the next one. # See /knockoff/datasets/__init__.py for mapping modelfamily = datasets.dataset_to_modelfamily[dataset_name] # e.g. 'classification' metadata = datasets.dataset_metadata[dataset_name] # Relevant parameters for the task. e.g. 'ngram' ngrams = metadata['ngram'] # If dataset does not exist, download it and save it dataset_dir = datadir + '/' + dataset_name.lower() + '_csv' train_data_path = os.path.join(dataset_dir, dataset_name + "_ngrams_{}_train.data".format(ngrams)) test_data_path = os.path.join(dataset_dir, dataset_name + "_ngrams_{}_test.data".format(ngrams)) if not (os.path.exists(train_data_path) and os.path.exists(test_data_path)): if not os.path.exists('.data'): print("Creating directory {}".format(datadir)) os.mkdir('.data') trainset, testset = text_classification.DATASETS[dataset_name](root='.data', ngrams=ngrams) print("Saving train data to {}".format(train_data_path)) torch.save(trainset, train_data_path) print("Saving test data to {}".format(test_data_path)) torch.save(testset, test_data_path) else: print("Loading train data from {}".format(train_data_path)) trainset = torch.load(train_data_path) print("Loading test data from {}".format(test_data_path)) testset = torch.load(test_data_path) # Extract variables for model from the dataset vocab_size = len(trainset.get_vocab()) params['num_classes'] = len(trainset.get_labels()) num_classes = params['num_classes'] model = zoo.get_net(model_name, modelfamily, pretrained, vocab_size=vocab_size, embed_dim=embed_dim, num_class=num_classes) model = model.to(device) model_utils.train_and_valid(trainset, testset, model, model_name, modelfamily, out_path, batch_size, lr, lr_gamma, num_workers, device=device, num_epochs=num_epochs) # Store arguments in json file. Maybe for the transfer set step? params['created_on'] = str(datetime.now()) params_out_path = osp.join(out_path, 'params.json') with open(params_out_path, 'w') as jf: json.dump(params, jf, indent=True)
def main(): parser = argparse.ArgumentParser(description='Train a model') # Required arguments parser.add_argument('model_dir', metavar='DIR', type=str, help='Directory containing transferset.pickle') parser.add_argument('model_arch', metavar='MODEL_ARCH', type=str, help='Model name') parser.add_argument('testdataset', metavar='DS_NAME', type=str, help='Name of test') parser.add_argument( '--budgets', metavar='B', type=str, help= 'Comma separated values of budgets. Knockoffs will be trained for each budget.' ) # Optional arguments parser.add_argument('-d', '--device_id', metavar='D', type=int, help='Device id. -1 for CPU.', default=0) parser.add_argument('-b', '--batch-size', type=int, default=64, metavar='N', help='input batch size for training (default: 64)') parser.add_argument('-e', '--epochs', type=int, default=100, metavar='N', help='number of epochs to train (default: 100)') parser.add_argument('--lr', type=float, default=0.01, metavar='LR', help='learning rate (default: 0.01)') parser.add_argument('--momentum', type=float, default=0.5, metavar='M', help='SGD momentum (default: 0.5)') parser.add_argument( '--log-interval', type=int, default=50, metavar='N', help='how many batches to wait before logging training status') parser.add_argument('--resume', default=None, type=str, metavar='PATH', help='path to latest checkpoint (default: none)') parser.add_argument('--lr-step', type=int, default=60, metavar='N', help='Step sizes for LR') parser.add_argument('--lr-gamma', type=float, default=0.1, metavar='N', help='LR Decay Rate') parser.add_argument('-w', '--num_workers', metavar='N', type=int, help='# Worker threads to load data', default=10) parser.add_argument('--pretrained', type=str, help='Use pretrained network', default=None) parser.add_argument('--weighted-loss', action='store_true', help='Use a weighted loss', default=False) args = parser.parse_args() params = vars(args) torch.manual_seed(cfg.DEFAULT_SEED) if params['device_id'] >= 0: os.environ["CUDA_VISIBLE_DEVICES"] = str(params['device_id']) device = torch.device('cuda') else: device = torch.device('cpu') model_dir = params['model_dir'] # ----------- Set up transferset transferset_path = osp.join(model_dir, 'transferset.pickle') with open(transferset_path, 'rb') as rf: transferset_samples = pickle.load(rf) num_classes = transferset_samples[0][1].size(0) print('=> found transfer set with {} samples, {} classes'.format( len(transferset_samples), num_classes)) # ----------- Set up testset dataset_name = params['testdataset'] valid_datasets = datasets.__dict__.keys() modelfamily = datasets.dataset_to_modelfamily[dataset_name] transform = datasets.modelfamily_to_transforms[modelfamily]['test'] if dataset_name not in valid_datasets: raise ValueError( 'Dataset not found. Valid arguments = {}'.format(valid_datasets)) dataset = datasets.__dict__[dataset_name] testset = dataset(train=False, transform=transform) if len(testset.classes) != num_classes: raise ValueError( '# Transfer classes ({}) != # Testset classes ({})'.format( num_classes, len(testset.classes))) # ----------- Set up model model_name = params['model_arch'] pretrained = params['pretrained'] # model = model_utils.get_net(model_name, n_output_classes=num_classes, pretrained=pretrained) model = zoo.get_net(model_name, modelfamily, pretrained, num_classes=num_classes) model = model.to(device) # ----------- Train budgets = [int(b) for b in params['budgets'].split(',')] for b in budgets: np.random.seed(cfg.DEFAULT_SEED) torch.manual_seed(cfg.DEFAULT_SEED) torch.cuda.manual_seed(cfg.DEFAULT_SEED) transferset = samples_to_transferset(transferset_samples, budget=b, transform=transform) print() print('=> Training at budget = {}'.format(len(transferset))) checkpoint_suffix = '.{}'.format(b) criterion_train = model_utils.soft_cross_entropy model_utils.train_model(model, transferset, model_dir, testset=testset, criterion_train=criterion_train, checkpoint_suffix=checkpoint_suffix, device=device, **params) # Store arguments params['created_on'] = str(datetime.now()) params_out_path = osp.join(model_dir, 'params_train.json') with open(params_out_path, 'w') as jf: json.dump(params, jf, indent=True)
def main(): parser = argparse.ArgumentParser( description='Construct apaptive transfer set') parser.add_argument( 'victim_model_dir', metavar='PATH', type=str, help= 'Path to victim model. Should contain files "model_best.pth.tar" and "params.json"' ) parser.add_argument('--out_dir', metavar='PATH', type=str, help='Destination directory to store transfer set', required=True) parser.add_argument('--budget', metavar='N', type=int, help='Size of transfer set to construct', required=True) parser.add_argument('model_arch', metavar='MODEL_ARCH', type=str, help='Model name') parser.add_argument('testdataset', metavar='DS_NAME', type=str, help='Name of test') parser.add_argument('--queryset', metavar='TYPE', type=str, help='Adversary\'s dataset (P_A(X))', required=True) parser.add_argument('--batch_size', metavar='TYPE', type=int, help='Batch size of queries', default=8) # parser.add_argument('--topk', metavar='N', type=int, help='Use posteriors only from topk classes', # default=None) # parser.add_argument('--rounding', metavar='N', type=int, help='Round posteriors to these many decimals', # default=None) # parser.add_argument('--tau_data', metavar='N', type=float, help='Frac. of data to sample from Adv data', # default=1.0) # parser.add_argument('--tau_classes', metavar='N', type=float, help='Frac. of classes to sample from Adv data', # default=1.0) # ----------- Other params parser.add_argument('-d', '--device_id', metavar='D', type=int, help='Device id', default=0) parser.add_argument('--pretrained', type=str, help='Use pretrained network', default=None) parser.add_argument('--defense', type=str, help='Defense strategy used by victim side', default=None) args = parser.parse_args() params = vars(args) out_path = params['out_dir'] knockoff_utils.create_dir(out_path + "-adaptive") defense = params['defense'] if defense: transfer_out_path = osp.join(out_path + "-adaptive", 'transferset-' + defense + '.pickle') else: transfer_out_path = osp.join(out_path + "-adaptive", 'transferset.pickle') torch.manual_seed(cfg.DEFAULT_SEED) if params['device_id'] >= 0: os.environ["CUDA_VISIBLE_DEVICES"] = str(params['device_id']) device = torch.device('cuda') else: device = torch.device('cpu') # ----------- Set up queryset queryset_name = params['queryset'] valid_datasets = datasets.__dict__.keys() if queryset_name not in valid_datasets: raise ValueError( 'Dataset not found. Valid arguments = {}'.format(valid_datasets)) modelfamily = datasets.dataset_to_modelfamily[queryset_name] transform = datasets.modelfamily_to_transforms[modelfamily]['test'] queryset = datasets.__dict__[queryset_name](train=True, transform=transform) queryset.targets = [ queryset.samples[idx][1] for idx in range(len(queryset)) ] # code.interact(local=dict(globals(), **locals())) num_classes = len(queryset.classes) # code.interact(local=dict(globals(), **locals())) # ----------- Initialize blackbox blackbox_dir = params['victim_model_dir'] blackbox = Blackbox.from_modeldir(blackbox_dir, defense, device) # ----------- Initialize Knockoff Nets model_name = params['model_arch'] pretrained = params['pretrained'] # model = model_utils.get_net(model_name, n_output_classes=num_classes, pretrained=pretrained) model = zoo.get_net(model_name, modelfamily, pretrained, num_classes=num_classes) model = model.to(device) adversary = AdaptiveAdversary(queryset, blackbox, model, device, reward='all') print('=> constructing transfer set...') transferset = adversary.get_transferset(params['budget']) with open(transfer_out_path, 'wb') as wf: pickle.dump(transferset, wf) print('=> transfer set ({} samples) written to: {}'.format( len(transferset), transfer_out_path)) # Store arguments params['created_on'] = str(datetime.now()) params_out_path = osp.join(out_path, 'params_transfer.json') with open(params_out_path, 'w') as jf: json.dump(params, jf, indent=True)
def main(): parser = argparse.ArgumentParser(description='Train a model') # Required arguments parser.add_argument('model_dir', metavar='DIR', type=str, help='Directory containing transferset.pickle') parser.add_argument('model_arch', metavar='MODEL_ARCH', type=str, help='Model name') parser.add_argument('testdataset', metavar='DS_NAME', type=str, help='Name of test') parser.add_argument('-o', '--output', type=str, help="Output dir") # Optional arguments parser.add_argument('-b', '--batch-size', type=int, default=1024, metavar='N', help='input batch size for training (default: 64)') parser.add_argument('--resume', default=None, type=str, metavar='PATH', help='path to latest checkpoint (default: none)') parser.add_argument('--device_id', default=0, type=int) parser.add_argument('-w', '--num_workers', metavar='N', type=int, help='# Worker threads to load data', default=10) args = parser.parse_args() params = vars(args) random.seed(cfg.DEFAULT_SEED) np.random.seed(cfg.DEFAULT_SEED) torch.cuda.manual_seed(cfg.DEFAULT_SEED) torch.manual_seed(cfg.DEFAULT_SEED) if params['device_id'] >= 0: os.environ["CUDA_VISIBLE_DEVICES"] = str(params['device_id']) device = torch.device('cuda') else: device = torch.device('cpu') model_dir = params['model_dir'] output_dir = params['output'] if not os.path.exists(output_dir): os.makedirs(output_dir) if params['testdataset'] == 'CIFAR10': num_classes = 10 # ----------- Set up testset dataset_name = params['testdataset'] valid_datasets = datasets.__dict__.keys() modelfamily = datasets.dataset_to_modelfamily[dataset_name] transform = datasets.modelfamily_to_transforms[modelfamily]['test'] if dataset_name not in valid_datasets: raise ValueError('Dataset not found. Valid arguments = {}'.format(valid_datasets)) dataset = datasets.__dict__[dataset_name] testset = dataset(train=False, transform=transform) if len(testset.classes) != num_classes: raise ValueError('# Transfer classes ({}) != # Testset classes ({})'.format(num_classes, len(testset.classes))) # ----------- Set up model model_name = params['model_arch'] # model = model_utils.get_net(model_name, n_output_classes=num_classes, pretrained=pretrained) model = zoo.get_net(model_name, modelfamily, False, num_classes=num_classes) # model = model.to(device) # model = protect_model(model) if params['model_arch'] == "resnet10": model = get_transferred_pytorch_model_resnet10() elif params['model_arch'] == "resnet18": model = get_transferred_pytorch_model_resnet18(testset) else: raise NotImplementedError model = model.cuda() # Test clean tf model accuracy # model = load_class_avg(model, arch=params['model_arch']) # test( # model=model, # testset=testset, # out_path=output_dir, # **params, # ) # st() # Test accuracy of each conv # for class_id in ["all"] + list(range(10)): # testset = dataset(train=False, transform=transform) # test_prune_per_conv( # model=model, # testset=testset, # out_path=output_dir, # class_id=class_id, # **params # ) # Test allconv pruned model om the whole dataset # test_prune_allconv_ratio_all( # model=model, # testset=testset, # out_path=output_dir, # **params # ) test_prune_allconv_ratio_subsetacc( model=model, testset=testset, out_path=output_dir, **params )
def __init__(self, blackbox, budget, model_adv_name, model_adv_pretrained, modelfamily, seedset, testset, device, out_dir, batch_size=cfg.DEFAULT_BATCH_SIZE, train_epochs=20, kappa=400, tau=None, rho=6, sigma=-1, query_batch_size=1, aug_strategy='jbda', useprobs=True, final_train_epochs=100): self.blackbox = blackbox self.budget = budget self.model_adv_name = model_adv_name self.model_adv_pretrained = model_adv_pretrained self.model_adv = None self.modelfamily = modelfamily self.seedset = seedset self.testset = testset self.batch_size = batch_size self.query_batch_size = query_batch_size self.testloader = DataLoader(self.testset, batch_size=self.batch_size, pin_memory=True) self.train_epochs = train_epochs self.final_train_epochs = final_train_epochs self.kappa = kappa self.tau = tau self.rho = rho self.sigma = sigma self.device = device self.out_dir = out_dir self.num_classes = len(self.testset.classes) assert (aug_strategy in ['jbda', 'jbself']) or 'jbtop' in aug_strategy self.aug_strategy = aug_strategy self.topk = 0 if 'jbtop' in aug_strategy: # extract k from "jbtop<k>" self.topk = int(aug_strategy.replace('jbtop', '')) self.accuracies = [] # Track test accuracies over time self.useprobs = useprobs # -------------------------- Initialize seed data print('=> Obtaining predictions over {} seed samples using strategy {}'.format(len(self.seedset), self.aug_strategy)) Dx = torch.cat([self.seedset[i][0].unsqueeze(0) for i in range(len(self.seedset))]) Dy = [] # Populate Dy with torch.no_grad(): for inputs, in DataLoader(TensorDataset(Dx), batch_size=self.query_batch_size): inputs = inputs.to(self.device) outputs = blackbox(inputs).cpu() if not self.useprobs: labels = torch.argmax(outputs, dim=1) labels_onehot = make_one_hot(labels, outputs.shape[1]) outputs = labels_onehot Dy.append(outputs) # Dy = torch.tensor(Dy) Dy = torch.cat(Dy) # TensorDataset D self.D = TensorDataset(Dx, Dy) ### Block memory required for training later on model_adv = zoo.get_net(self.model_adv_name, self.modelfamily, self.model_adv_pretrained, num_classes=self.num_classes) model_adv = model_adv.to(self.device) model_adv = model_utils.train_model(model_adv, self.D, self.out_dir, num_workers=10, checkpoint_suffix='.{}'.format(self.blackbox.call_count), device=self.device, epochs=1, log_interval=500, lr=0.01, momentum=0.9, batch_size=self.batch_size, lr_gamma=0.1, testset=self.testset, criterion_train=model_utils.soft_cross_entropy)
def get_transferred_pytorch_model_resnet10(): model = zoo.get_net("resnet10", "cifar", False, num_classes=10) model = load_tf_model_resnet10(model) return model
def main(): parser = argparse.ArgumentParser(description='Train a model') # Required arguments parser.add_argument('policy', metavar='PI', type=str, help='Policy to use while training', choices=['random', 'adaptive', 'jacobian']) parser.add_argument('model_dir', metavar='DIR', type=str, help='Directory containing transferset.pickle') parser.add_argument('model_arch', metavar='MODEL_ARCH', type=str, help='Model name') parser.add_argument('testdataset', metavar='DS_NAME', type=str, help='Name of test') parser.add_argument( '--budgets', metavar='B', type=str, help= 'Comma separated values of budgets. Knockoffs will be trained for each budget.' ) # Optional arguments parser.add_argument('-d', '--device_id', metavar='D', type=int, help='Device id. -1 for CPU.', default=0) parser.add_argument('-b', '--batch-size', type=int, default=64, metavar='N', help='input batch size for training (default: 64)') parser.add_argument('-e', '--epochs', type=int, default=100, metavar='N', help='number of epochs to train (default: 100)') parser.add_argument('--lr', type=float, default=0.1, metavar='LR', help='learning rate (default: 0.01)') parser.add_argument('--momentum', type=float, default=0.5, metavar='M', help='SGD momentum (default: 0.5)') parser.add_argument( '--log-interval', type=int, default=50, metavar='N', help='how many batches to wait before logging training status') parser.add_argument('--resume', default=None, type=str, metavar='PATH', help='path to latest checkpoint (default: none)') parser.add_argument('--lr-step', type=int, default=60, metavar='N', help='Step sizes for LR') parser.add_argument('--lr-gamma', type=float, default=0.1, metavar='N', help='LR Decay Rate') parser.add_argument('-w', '--num_workers', metavar='N', type=int, help='# Worker threads to load data', default=10) parser.add_argument('--pretrained', type=str, help='Use pretrained network', default=None) parser.add_argument('--weighted-loss', action='store_true', help='Use a weighted loss', default=False) # Attacker's defense parser.add_argument('--argmaxed', action='store_true', help='Only consider argmax labels', default=False) parser.add_argument('--optimizer_choice', type=str, help='Optimizer', default='sgdm', choices=('sgd', 'sgdm', 'adam', 'adagrad')) parser.add_argument('--defense', type=str, help='Defense strategy used by victim side', default=None) args = parser.parse_args() params = vars(args) torch.manual_seed(cfg.DEFAULT_SEED) if params['device_id'] >= 0: os.environ["CUDA_VISIBLE_DEVICES"] = str(params['device_id']) device = torch.device('cuda') else: device = torch.device('cpu') model_dir = params['model_dir'] + "-" + params["policy"] # ----------- Set up transferset defense = params['defense'] if defense: transferset_path = osp.join(model_dir, 'transferset-' + defense + '.pickle') else: transferset_path = osp.join(model_dir, 'transferset.pickle') with open(transferset_path, 'rb') as rf: transferset_samples = pickle.load(rf) num_classes = transferset_samples[0][1].size(0) print('=> found transfer set with {} samples, {} classes'.format( len(transferset_samples), num_classes)) # ----------- Clean up transfer (if necessary) if params['argmaxed']: new_transferset_samples = [] print('=> Using argmax labels (instead of posterior probabilities)') for i in range(len(transferset_samples)): x_i, y_i = transferset_samples[i] argmax_k = y_i.argmax() y_i_1hot = torch.zeros_like(y_i) y_i_1hot[argmax_k] = 1. new_transferset_samples.append((x_i, y_i_1hot)) transferset_samples = new_transferset_samples # ----------- Set up testset dataset_name = params['testdataset'] valid_datasets = datasets.__dict__.keys() modelfamily = datasets.dataset_to_modelfamily[dataset_name] transform = datasets.modelfamily_to_transforms[modelfamily]['test'] if dataset_name not in valid_datasets: raise ValueError( 'Dataset not found. Valid arguments = {}'.format(valid_datasets)) dataset = datasets.__dict__[dataset_name] testset = dataset(train=False, transform=transform) if len(testset.classes) != num_classes: raise ValueError( '# Transfer classes ({}) != # Testset classes ({})'.format( num_classes, len(testset.classes))) # ----------- Set up model model_name = params['model_arch'] pretrained = params['pretrained'] # model = model_utils.get_net(model_name, n_output_classes=num_classes, pretrained=pretrained) model = zoo.get_net(model_name, modelfamily, pretrained, num_classes=num_classes) model = model.to(device) # ----------- Train budgets = [int(b) for b in params['budgets'].split(',')] for b in budgets: np.random.seed(cfg.DEFAULT_SEED) torch.manual_seed(cfg.DEFAULT_SEED) torch.cuda.manual_seed(cfg.DEFAULT_SEED) transferset = samples_to_transferset(transferset_samples, budget=b, transform=transform) # code.interact(local=dict(globals(), **locals())) print() print('=> Training at budget = {}'.format(len(transferset))) optimizer = get_optimizer(model.parameters(), params['optimizer_choice'], **params) print(params) checkpoint_suffix = '.{}'.format(b) criterion_train = model_utils.soft_cross_entropy model_utils.train_model(model, transferset, model_dir, testset=testset, criterion_train=criterion_train, checkpoint_suffix=checkpoint_suffix, device=device, optimizer=optimizer, **params) # Store arguments params['created_on'] = str(datetime.now()) params_out_path = osp.join(model_dir, 'params_train.json') with open(params_out_path, 'w') as jf: json.dump(params, jf, indent=True)
def get_transferred_pytorch_model_resnet18(testset): model = zoo.get_net("resnet18", "cifar", False, num_classes=10) model = load_tf_model_resnet18(model, testset) return model
def main(): parser = argparse.ArgumentParser(description='Train a model') # Required arguments parser.add_argument('model_dir', metavar='DIR', type=str, help='Directory containing transferset.pickle') parser.add_argument('model_arch', metavar='MODEL_ARCH', type=str, help='Model name') parser.add_argument('testdataset', metavar='DS_NAME', type=str, help='Name of test') parser.add_argument('--budgets', metavar='B', type=int, help='Knockoffs will be trained for budget.') # Optional arguments parser.add_argument('-d', '--device_id', metavar='D', type=int, help='Device id. -1 for CPU.', default=0) parser.add_argument('-b', '--batch-size', type=int, default=64, metavar='N', help='input batch size for training (default: 64)') parser.add_argument('-e', '--epochs', type=int, default=10, metavar='N', help='number of epochs to train (default: 10)') parser.add_argument('--lr', type=float, default=0.01, metavar='LR', help='learning rate (default: 0.01)') parser.add_argument('--momentum', type=float, default=0.5, metavar='M', help='SGD momentum (default: 0.5)') parser.add_argument( '--log-interval', type=int, default=50, metavar='N', help='how many batches to wait before logging training status') parser.add_argument('--resume', default=None, type=str, metavar='PATH', help='path to latest checkpoint (default: none)') parser.add_argument('--lr-step', type=int, default=60, metavar='N', help='Step sizes for LR') parser.add_argument('--lr-gamma', type=float, default=0.1, metavar='N', help='LR Decay Rate') parser.add_argument('-w', '--num_workers', metavar='N', type=int, help='# Worker threads to load data', default=10) parser.add_argument('--pretrained', type=str, help='Use pretrained network', default=None) parser.add_argument('--weighted-loss', action='store_true', help='Use a weighted loss', default=False) # RL arguments parser.add_argument('--traj_length', metavar='N', type=int, help='# Step in one trajactory', default=10) parser.add_argument('--num_each_class', metavar='N', type=int, help='# sample in each class', default=1) parser.add_argument('--n_iter', metavar='N', type=int, help='# iterations of RL training', default=10) parser.add_argument('--n_traj_each_iter', metavar='N', type=int, help='# trajactories / iter', default=10) parser.add_argument('--queryset', metavar='DS_NAME', type=str, help='Name of test') parser.add_argument('--victim_model_dir', default=None, type=str, metavar='PATH', help='path to latest checkpoint (default: none)') parser.add_argument('--n_layers', metavar='N', type=int, help='# layers in policy', default=4) parser.add_argument('--size', metavar='N', type=int, help='size of layer in policy', default=64) parser.add_argument('--policy_lr', type=float, default=1e-4, metavar='N', help='Policy learning rate') parser.add_argument('--num_agent_train_steps_per_iter', metavar='N', type=int, help='num_agent_train_steps_per_iter', default=10) parser.add_argument('--agent_train_batch_size', metavar='N', type=int, help='num_agent_train_steps_per_iter', default=990) parser.add_argument('--policy_gamma', type=float, default=0.9, metavar='N', help='reward discounting') parser.add_argument('--eps_random', type=float, default=-1, metavar='N', help='eps random exploration') parser.add_argument('--nn_baseline', action='store_true', help='Use nn baseline', default=False) # Attacker's defense parser.add_argument('--argmaxed', action='store_true', help='Only consider argmax labels', default=False) parser.add_argument('--optimizer_choice', type=str, help='Optimizer', default='sgdm', choices=('sgd', 'sgdm', 'adam', 'adagrad')) args = parser.parse_args() params = vars(args) torch.manual_seed(cfg.DEFAULT_SEED) if params['device_id'] >= 0: os.environ["CUDA_VISIBLE_DEVICES"] = str(params['device_id']) device = torch.device('cuda') ptu.init_gpu() else: device = torch.device('cpu') model_dir = params['model_dir'] # ----------- Set up testset dataset_name = params['testdataset'] valid_datasets = datasets.__dict__.keys() modelfamily = datasets.dataset_to_modelfamily[dataset_name] transform = datasets.modelfamily_to_transforms[modelfamily]['test'] if dataset_name not in valid_datasets: raise ValueError( 'Dataset not found. Valid arguments = {}'.format(valid_datasets)) dataset = datasets.__dict__[dataset_name] testset = dataset(train=False, transform=transform) #if len(testset.classes) != num_classes: # raise ValueError('# Transfer classes ({}) != # Testset classes ({})'.format(num_classes, len(testset.classes))) # ----------- Set up queryset queryset_name = params['queryset'] valid_datasets = datasets.__dict__.keys() if queryset_name not in valid_datasets: raise ValueError( 'Dataset not found. Valid arguments = {}'.format(valid_datasets)) modelfamily = datasets.dataset_to_modelfamily[queryset_name] transform = datasets.modelfamily_to_transforms[modelfamily]['train'] try: queryset = datasets.__dict__[queryset_name](train=True, transform=transform) except: queryset = datasets.__dict__[queryset_name](split="train", transform=transform) num_classes = len(queryset.classes) # ----------- Initialize blackbox blackbox_dir = params['victim_model_dir'] blackbox = Blackbox.from_modeldir(blackbox_dir, device) # ----------- Set up adversary model model_name = params['model_arch'] pretrained = params['pretrained'] # model = model_utils.get_net(model_name, n_output_classes=num_classes, pretrained=pretrained) adv_model = zoo.get_net(model_name, modelfamily, pretrained, num_classes=10) adv_model = adv_model.to(device) # ----------- Initialize adversary num_each_class = params['num_each_class'] agent_params = { "ac_dim": num_classes, "ob_dim": len(testset.classes), "n_layers": params["n_layers"], "size": params["size"], "discrete": True, "learning_rate": params["policy_lr"], "num_agent_train_steps_per_iter": params["num_agent_train_steps_per_iter"], "agent_train_batch_size": params["agent_train_batch_size"], "gamma": params["policy_gamma"], "reward_to_go": True, "nn_baseline": params["nn_baseline"], "standardize_advantages": True, "eps_random": params["eps_random"] } adversary = PGAdversary(queryset, num_each_class, agent_params) # ----------- Set up transferset def collect_training_trajectories(length, n_traj=10): nonlocal avg_rewards, avg_components paths = [] mean_rew = 0 mean_cert = mean_L = mean_E = mean_div = 0 X_paths, Y_paths = [], [] for _ in range(n_traj): obs, acs, rewards, next_obs = [], [], [], [] r_certs, r_Ls, r_Es, r_divs = [], [], [], [] X_path, Y_path = [], [] X, actions = adversary.init_sampling() X_path.append(X) ob = blackbox(X) Y_path.append(ob) ob = ob.numpy() for t in range(length - 1): with torch.no_grad(): # Observe and react obs.append(ob) X_new, actions = adversary.sample(ob) X_path.append(X_new) acs.append(actions) # Env gives feedback, which is a new observation X_new = X_new.to(device) ob = blackbox(X_new) Y_path.append(ob) ob = ob.cpu().numpy() next_obs.append(ob) Y_adv = adv_model(X_new) Y_adv = F.softmax(Y_adv, dim=1).cpu().numpy() reward, r_cert, r_L, r_E, r_div = adversary.agent.calculate_reward( ob, np.concatenate(acs), Y_adv) rewards.append(reward) r_certs.append(r_cert) r_Ls.append(r_L) r_Es.append(r_E) r_divs.append(r_div) obs = np.concatenate(obs) acs = np.concatenate(acs) rewards = np.concatenate(rewards) mean_rew += np.mean(rewards) mean_cert += np.mean(np.concatenate(r_certs)) mean_L += np.mean(np.concatenate(r_Ls)) mean_E += np.mean(np.array(r_Es)) mean_div += np.mean(np.array(r_divs)) next_obs = np.concatenate(next_obs) path = { "observation": obs, "action": acs, "reward": rewards, "next_observation": next_obs } paths.append(path) X_paths.append(torch.cat(X_path)) Y_paths.append(torch.cat(Y_path)) print(f"==> Avg reward: {mean_rew / n_traj}") avg_rewards.append(mean_rew / n_traj) avg_components["avg_cert"].append(mean_cert / n_traj) avg_components["avg_L"].append(mean_L / n_traj) avg_components["avg_E"].append(mean_E / n_traj) avg_components["avg_div"].append(mean_div / n_traj) return torch.cat(X_paths), torch.cat(Y_paths), paths traj_length = params['traj_length'] num_each_class = params['num_each_class'] n_iter = params['n_iter'] X, Y = None, None budgets = params['budgets'] n_traj = params['n_traj_each_iter'] criterion_train = model_utils.soft_cross_entropy if traj_length > 0: n_iter = budgets // (traj_length * n_traj) print(f"==> Budget = {n_iter} x {traj_length} x {n_traj}") best_test_acc = [] best_acc = -1 avg_rewards = [] avg_components = collections.defaultdict(list) for iter in range(1, n_iter + 1): # n_iter * traj_length = budget print(f"==> Iteration: {iter}/{n_iter}") X_path, Y_path, paths = collect_training_trajectories(traj_length, n_traj=n_traj) adversary.add_to_replay_buffer(paths) adversary.train_agent() if X is None: X, Y = X_path, Y_path else: X = torch.cat((X, X_path)) Y = torch.cat((Y, Y_path)) transferset = ImageTensorSet((X, Y)) # ----------- Train #np.random.seed(cfg.DEFAULT_SEED) #torch.manual_seed(cfg.DEFAULT_SEED) #torch.cuda.manual_seed(cfg.DEFAULT_SEED) optimizer = get_optimizer(adv_model.parameters(), params['optimizer_choice'], **params) print(f"Train on {len(transferset)} samples") checkpoint_suffix = '.{extraction}' best_acc = model_utils.train_model(adv_model, transferset, model_dir, testset=testset, criterion_train=criterion_train, checkpoint_suffix=checkpoint_suffix, device=device, optimizer=optimizer, benchmark=best_acc, **params) best_test_acc.append(best_acc) adversary.agent.actor.save( osp.join(model_dir, "checkpoint.agent.state_dict")) # ----------- Log torch.save(best_test_acc, osp.join(model_dir, "best_acc.pylist")) torch.save(avg_rewards, osp.join(model_dir, "avg_rewards.pylist")) torch.save(avg_components, osp.join(model_dir, "avg_components.pydict")) torch.save(adversary.idx_counter, osp.join(model_dir, "idx_counter.pydict")) torch.save(transferset, osp.join(model_dir, "transferset.pt")) # Store arguments params['created_on'] = str(datetime.now()) params_out_path = osp.join(model_dir, 'params_train.json') with open(params_out_path, 'w') as jf: json.dump(params, jf, indent=True) agent_params_out_path = osp.join(model_dir, 'agent_params_train.json') with open(agent_params_out_path, 'w') as jf: json.dump(agent_params, jf, indent=True)