def __init__(self, model_ori, pred, test, solve_slope=False, device='cuda', simplify=True, in_size=(1, 3, 32, 32)): """ convert pytorch model to auto_LiRPA module """ layers = list(model_ori.children()) if simplify: added_prop_layers = add_single_prop(layers, pred, test) self.layers = added_prop_layers else: self.layers = layers net = nn.Sequential(*self.layers) self.solve_slope = solve_slope if solve_slope: self.net = BoundedModule(net, torch.rand(in_size), bound_opts={ 'relu': 'random_evaluation', 'conv_mode': 'patches' }, device=device) else: self.net = BoundedModule(net, torch.rand(in_size), bound_opts={'relu': 'same-slope'}, device=device) self.net.eval()
def __init__(self, state_dim, action_dim, actor_network, critic_network, mini_batch_size, actor_opt_fn, critic_opt_fn, robust_params=None): super(RobustDeterministicActorCriticNet, self).__init__() if robust_params is None: robust_params = {} self.use_loss_fusion = robust_params.get('use_loss_fusion', False) # Use loss fusion to reduce complexity for convex relaxation. Default is False. self.use_full_backward = robust_params.get('use_full_backward', False) if self.use_loss_fusion: # Use auto_LiRPA to compute the L2 norm directly. self.fc_action = model_mlp_any_with_loss(state_dim, actor_network, action_dim) modules = self.fc_action._modules # Auto LiRPA wrapper self.fc_action = BoundedModule( self.fc_action, (torch.empty(size=(1, state_dim)), torch.empty(size=(1, action_dim))), device=Config.DEVICE) # self.fc_action._modules = modules for n in self.fc_action.nodes: # Find the tanh neuron in computational graph if isinstance(n, BoundTanh): self.fc_action_after_tanh = n self.fc_action_pre_tanh = n.inputs[0] break else: # Fully connected layer with [state_dim, 400, 300, action_dim] neurons and ReLU activation function self.fc_action = model_mlp_any(state_dim, actor_network, action_dim) # auto_lirpa wrapper self.fc_action = BoundedModule( self.fc_action, (torch.empty(size=(1, state_dim)), ), device=Config.DEVICE) # Fully connected layer with [state_dim + action_dim, 400, 300, 1] self.fc_critic = model_mlp_any(state_dim + action_dim, critic_network, 1) # auto_lirpa wrapper self.fc_critic = BoundedModule( self.fc_critic, (torch.empty(size=(1, state_dim + action_dim)), ), device=Config.DEVICE) self.actor_params = self.fc_action.parameters() self.critic_params = self.fc_critic.parameters() self.actor_opt = actor_opt_fn(self.actor_params) self.critic_opt = critic_opt_fn(self.critic_params) self.to(Config.DEVICE) # Create identity specification matrices self.actor_identity = torch.eye(action_dim).repeat(mini_batch_size,1,1).to(Config.DEVICE) self.critic_identity = torch.eye(1).repeat(mini_batch_size,1,1).to(Config.DEVICE) self.action_dim = action_dim self.state_dim = state_dim
def verify_model(pert_weight=True, pert_bias=True, norm=np.inf, lb_name='', ub_name=''): model_ori_ = mlp_3layer(pert_weight=pert_weight, pert_bias=pert_bias, norm=norm).eval() model_ori_.load_state_dict(state_dict) model_ = BoundedModule(model_ori_, inputs) ptb.eps = 0.01 verify_bounds(model_, dummy_input, IBP=True, method='backward', forward_ret=forward_ret, lb_name=lb_name + '_CROWN-IBP', ub_name=ub_name + '_CROWN-IBP') # CROWN-IBP verify_bounds(model_, dummy_input, IBP=False, method='backward', forward_ret=forward_ret, lb_name=lb_name + '_CROWN', ub_name=ub_name + '_CROWN') # CROWN
def test(self): model_oris = [ models.model_resnet(width=1, mult=2), models.ResNet18(in_planes=2) ] self.result = [] for model_ori in model_oris: conv_mode = 'patches' # conv_mode can be set as 'matrix' or 'patches' normalize = torchvision.transforms.Normalize( mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010]) test_data = torchvision.datasets.CIFAR10( "./data", train=False, download=True, transform=torchvision.transforms.Compose( [torchvision.transforms.ToTensor(), normalize])) N = 1 n_classes = 10 image = torch.Tensor(test_data.data[:N]).reshape(N, 3, 32, 32) image = image.to(torch.float32) / 255.0 model = BoundedModule(model_ori, image, bound_opts={"conv_mode": conv_mode}) ptb = PerturbationLpNorm(norm=np.inf, eps=0.03) image = BoundedTensor(image, ptb) pred = model(image) lb, ub = model.compute_bounds(IBP=False, C=None, method='backward') self.result += [lb, ub] self.check()
def test(): torch.manual_seed(1) torch.cuda.manual_seed_all(1) model = cnn_MNIST() checkpoint = torch.load("../examples/vision/pretrain/mnist_cnn_small.pth", map_location="cpu") model.load_state_dict(checkpoint) N = 2 n_classes = 10 image = torch.randn(N, 1, 28, 28) image = image.to(torch.float32) / 255.0 model = BoundedModule(model, torch.empty_like(image), device="cpu") eps = 0.3 norm = np.inf ptb = PerturbationLpNorm(norm=norm, eps=eps) image = BoundedTensor(image, ptb) pred = model(image) lb, ub = model.compute_bounds() assert lb.shape == ub.shape == torch.Size((2, 10)) path = 'data/constant_test_data' if args.gen_ref: torch.save((lb, ub), path) else: lb_ref, ub_ref = torch.load(path) print(lb) print(lb_ref) assert torch.allclose(lb, lb_ref) assert torch.allclose(ub, ub_ref)
def compute_and_compare_bounds(self, eps, norm, IBP, method): input_data = torch.randn((N, 256)) model = BoundedModule(self.original_model, torch.empty_like(input_data)) ptb = PerturbationLpNorm(norm=norm, eps=eps) ptb_data = BoundedTensor(input_data, ptb) pred = model(ptb_data) label = torch.argmax(pred, dim=1).cpu().detach().numpy() # Compute bounds. lb, ub = model.compute_bounds(IBP=IBP, method=method) # Compute dual norm. if norm == 1: q = np.inf elif norm == np.inf: q = 1.0 else: q = 1.0 / (1.0 - (1.0 / norm)) # Compute reference manually. weight, bias = list(model.parameters()) norm = weight.norm(p=q, dim=1) expected_pred = input_data.matmul(weight.t()) + bias expected_ub = eps * norm + expected_pred expected_lb = -eps * norm + expected_pred # Check equivalence. self.assertEqual(expected_pred, pred) self.assertEqual(expected_ub, ub) self.assertEqual(expected_lb, lb)
def train(epoch, batches, type): meter = MultiAverageMeter() assert(optimizer is not None) train = type == 'train' if args.robust: eps_scheduler.set_epoch_length(len(batches)) if train: eps_scheduler.train() eps_scheduler.step_epoch() else: eps_scheduler.eval() for i, batch in enumerate(batches): if args.robust: eps_scheduler.step_batch() eps = eps_scheduler.get_eps() else: eps = 0 acc, loss, acc_robust, loss_robust = \ step(model, ptb, batch, eps=eps, train=train) meter.update('acc', acc, len(batch)) meter.update('loss', loss, len(batch)) meter.update('acc_rob', acc_robust, len(batch)) meter.update('loss_rob', loss_robust, len(batch)) if train: if (i + 1) % args.gradient_accumulation_steps == 0 or (i + 1) == len(batches): scale_gradients(optimizer, i % args.gradient_accumulation_steps + 1, args.grad_clip) optimizer.step() optimizer.zero_grad() if lr_scheduler is not None: lr_scheduler.step() writer.add_scalar('loss_train_{}'.format(epoch), meter.avg('loss'), i + 1) writer.add_scalar('loss_robust_train_{}'.format(epoch), meter.avg('loss_rob'), i + 1) writer.add_scalar('acc_train_{}'.format(epoch), meter.avg('acc'), i + 1) writer.add_scalar('acc_robust_train_{}'.format(epoch), meter.avg('acc_rob'), i + 1) if (i + 1) % args.log_interval == 0 or (i + 1) == len(batches): logger.info('Epoch {}, {} step {}/{}: eps {:.5f}, {}'.format( epoch, type, i + 1, len(batches), eps, meter)) if lr_scheduler is not None: logger.info('lr {}'.format(lr_scheduler.get_lr())) writer.add_scalar('loss/{}'.format(type), meter.avg('loss'), epoch) writer.add_scalar('loss_robust/{}'.format(type), meter.avg('loss_rob'), epoch) writer.add_scalar('acc/{}'.format(type), meter.avg('acc'), epoch) writer.add_scalar('acc_robust/{}'.format(type), meter.avg('acc_rob'), epoch) if train: if args.loss_fusion: state_dict_loss = model_loss.state_dict() state_dict = {} for name in state_dict_loss: assert(name.startswith('model.')) state_dict[name[6:]] = state_dict_loss[name] model_ori.load_state_dict(state_dict) model_bound = BoundedModule( model_ori, (dummy_embeddings, dummy_mask), bound_opts=bound_opts, device=args.device) model.model_from_embeddings = model_bound model.save(epoch) return meter.avg('acc_rob')
def setup_sarsa(self, lr_schedule, eps_scheduler, beta_scheduler): # Create the Sarsa model, with S and A as the input. self.sarsa_model = ValueDenseNet(self.NUM_FEATURES + self.NUM_ACTIONS, self.INITIALIZATION) self.sarsa_opt = optim.Adam(self.sarsa_model.parameters(), lr=self.VAL_LR, eps=1e-5) self.sarsa_scheduler = optim.lr_scheduler.LambdaLR(self.sarsa_opt, lr_schedule) self.sarsa_eps_scheduler = eps_scheduler self.sarsa_beta_scheduler = beta_scheduler # Convert model with relaxation wrapper. dummy_input = torch.randn(1, self.NUM_FEATURES + self.NUM_ACTIONS) self.relaxed_sarsa_model = BoundedModule(self.sarsa_model, dummy_input)
def test(): torch.manual_seed(1) torch.cuda.manual_seed_all(1) models = [2, 3] paddings = [1, 2] strides = [1, 3] N = 2 n_classes = 10 image = torch.randn(N, 1, 28, 28) image = image.to(torch.float32) / 255.0 for layer_num in models: for padding in paddings: for stride in strides: # print(layer_num, padding, stride) try: model_ori = cnn_model(layer_num, padding, stride) except: continue model = BoundedModule(model_ori, torch.empty_like(image), device="cpu", bound_opts={"conv_mode": "patches"}) eps = 0.3 norm = np.inf ptb = PerturbationLpNorm(norm=norm, eps=eps) image = BoundedTensor(image, ptb) pred = model(image) lb, ub = model.compute_bounds() model = BoundedModule(model_ori, torch.empty_like(image), device="cpu", bound_opts={"conv_mode": "matrix"}) pred = model(image) lb_ref, ub_ref = model.compute_bounds() assert lb.shape == ub.shape == torch.Size((N, n_classes)) assert torch.allclose(lb, lb_ref) assert torch.allclose(ub, ub_ref)
def torch2network(self, torch_model): from auto_LiRPA import BoundedModule my_input = torch.empty((1, ) + self.input_shape) if hasattr(torch_model, "core"): torch_model = torch_model.core model = BoundedModule(torch_model, my_input, bound_opts=self.bound_opts) return model
def test(): net = ResNet18() N = 2 n_classes = 10 x = torch.randn(N, 3, 32, 32) y = net(x) device = 'cpu' if device == 'cuda': x = x.cuda() y = y.cuda() model = BoundedModule(net, torch.empty_like(x), bound_opts={"conv_mode": "patches"}, device=device) print("Model structure: \n", str(net)) eps = 0.3 norm = np.inf ptb = PerturbationLpNorm(norm=norm, eps=eps) image = BoundedTensor(x, ptb) pred = model(image) lb, ub = model.compute_bounds() model = BoundedModule(net, torch.empty_like(x), bound_opts={"conv_mode": "matrix"}, device=device) eps = 0.3 norm = np.inf ptb = PerturbationLpNorm(norm=norm, eps=eps) image = BoundedTensor(x, ptb) pred = model(image) lb_ref, ub_ref = model.compute_bounds() # assert lb.shape == ub.shape == torch.Size((N, n_classes)) print((lb - lb_ref).sum(), (ub - ub_ref).sum()) assert torch.allclose(lb, lb_ref) assert torch.allclose(ub, ub_ref)
def create_relaxed_model(self, time_in_state=False): # Create state perturbation model for robust PPO training. if isinstance(self.policy_model, CtsPolicy): from .convex_relaxation import RelaxedCtsPolicyForState relaxed_policy_model = RelaxedCtsPolicyForState( self.NUM_FEATURES, self.NUM_ACTIONS, time_in_state=time_in_state, activation=self.policy_activation, policy_model=self.policy_model) dummy_input1 = torch.randn(1, self.NUM_FEATURES) inputs = (dummy_input1, ) self.relaxed_policy_model = BoundedModule(relaxed_policy_model, inputs) self.robust_eps_scheduler = LinearScheduler(self.params.ROBUST_PPO_EPS, self.params.ROBUST_PPO_EPS_SCHEDULER_OPTS) if self.params.ROBUST_PPO_BETA_SCHEDULER_OPTS == "same": self.robust_beta_scheduler = LinearScheduler(self.params.ROBUST_PPO_BETA, self.params.ROBUST_PPO_EPS_SCHEDULER_OPTS) else: self.robust_beta_scheduler = LinearScheduler(self.params.ROBUST_PPO_BETA, self.params.ROBUST_PPO_BETA_SCHEDULER_OPTS) else: raise NotImplementedError
def compute_and_compare_bounds(self, eps, norm, IBP, method): input_data = torch.randn((N, 1, input_dim, input_dim)) model = BoundedModule(self.original_model, torch.empty_like(input_data)) ptb = PerturbationLpNorm(norm=norm, eps=eps) ptb_data = BoundedTensor(input_data, ptb) pred = model(ptb_data) label = torch.argmax(pred, dim=1).cpu().detach().numpy() # Compute bounds. lb, ub = model.compute_bounds(IBP=IBP, method=method) # Compute reference. conv_weight, conv_bias = list(model.parameters()) conv_bias = conv_bias.view(1, out_channel, 1, 1) matrix_eye = torch.eye(input_dim * input_dim).view( input_dim * input_dim, 1, input_dim, input_dim) # Obtain equivalent weight and bias for convolution. weight = self.original_model.conv( matrix_eye ) - conv_bias # Output is (batch, channel, weight, height). weight = weight.view( input_dim * input_dim, -1) # Dimension is (flattened_input, flattened_output). bias = conv_bias.repeat(1, 1, input_dim // 2, input_dim // 2).view(-1) flattend_data = input_data.view(N, -1) # Compute dual norm. if norm == 1: q = np.inf elif norm == np.inf: q = 1.0 else: q = 1.0 / (1.0 - (1.0 / norm)) # Manually compute bounds. norm = weight.t().norm(p=q, dim=1) expected_pred = flattend_data.matmul(weight) + bias expected_ub = eps * norm + expected_pred expected_lb = -eps * norm + expected_pred # Check equivalence. if method == 'backward' or method == 'forward': self.assertEqual(expected_pred, pred) self.assertEqual(expected_ub, ub) self.assertEqual(expected_lb, lb)
def test_perturbation(self): np.random.seed(123) # FIXME This seed is inconsistent with other seeds (1234) model_ori = models.Models['mlp_3layer_weight_perturb'](pert_weight=True, pert_bias=True).eval() self.result['model'] = model_ori.state_dict() self.result['data'] = torch.randn(8, 1, 28, 28) model_ori.load_state_dict(self.result['model']) state_dict = copy.deepcopy(model_ori.state_dict()) dummy_input = self.result['data'].requires_grad_() inputs = (dummy_input,) model = BoundedModule(model_ori, inputs) forward_ret = model(dummy_input) model_ori.eval() assert torch.isclose(model_ori(dummy_input), model_ori(dummy_input), 1e-8).all() def verify_model(pert_weight=True, pert_bias=True, norm=np.inf, lb_name='', ub_name=''): model_ori_ = models.Models['mlp_3layer_weight_perturb'](pert_weight=pert_weight, pert_bias=pert_bias, norm=norm).eval() model_ori_.load_state_dict(state_dict) model_ = BoundedModule(model_ori_, inputs) model_.ptb = model_ori.ptb self.verify_bounds(model_, dummy_input, IBP=True, method='backward', forward_ret=forward_ret, lb_name=lb_name + '_CROWN-IBP', ub_name=ub_name + '_CROWN-IBP') # CROWN-IBP self.verify_bounds(model_, dummy_input, IBP=False, method='backward', forward_ret=forward_ret, lb_name=lb_name + '_CROWN', ub_name=ub_name + '_CROWN') # CROWN # Linf verify_model(pert_weight=True, pert_bias=True, norm=np.inf, lb_name='l_inf_weights_bias_lb', ub_name='l_inf_weights_bias_ub') verify_model(pert_weight=True, pert_bias=False, norm=np.inf, lb_name='l_inf_weights_lb', ub_name='l_inf_weights_ub') verify_model(pert_weight=False, pert_bias=True, norm=np.inf, lb_name='l_inf_bias_lb', ub_name='l_inf_bias_ub') # L2 verify_model(pert_weight=True, pert_bias=True, norm=2, lb_name='l_2_weights_bias_lb', ub_name='l_2_weights_bias_ub') verify_model(pert_weight=True, pert_bias=False, norm=2, lb_name='l_2_weights_lb', ub_name='l_2_weights_ub') verify_model(pert_weight=False, pert_bias=True, norm=2, lb_name='l_2_bias_lb', ub_name='l_2_bias_ub') if self.generate: self.save()
def __init__(self, name, env, input_shape, num_actions, robust=False, width=1): super(QNetwork, self).__init__() self.env = env self.input_shape = input_shape self.num_actions = num_actions self.robust = robust if name == 'DQN': self.features = nn.Sequential( nn.Linear(input_shape[0], 128 * width), nn.ReLU(), nn.Linear(128 * width, 128 * width), nn.ReLU(), nn.Linear(128 * width, self.env.action_space.n)) elif name == 'CnnDQN': self.features = nn.Sequential( nn.Conv2d(input_shape[0], 32 * width, kernel_size=8, stride=4), nn.ReLU(), nn.Conv2d(32 * width, 64 * width, kernel_size=4, stride=2), nn.ReLU(), nn.Conv2d(64 * width, 64 * width, kernel_size=3, stride=1), nn.ReLU(), Flatten(), nn.Linear(3136 * width, 512 * width), nn.ReLU(), nn.Linear(512 * width, self.num_actions)) elif name == 'DuelingCnnDQN': self.features = DuelingCnnDQN(input_shape, num_actions, width) else: raise NotImplementedError( '{} network structure not implemented.'.format(name)) if self.robust: dummy_input = torch.empty_like(torch.randn((1, ) + input_shape)) self.features = BoundedModule(self.features, dummy_input, device="cuda" if USE_CUDA else "cpu")
def test(self): model = cnn_MNIST() checkpoint = torch.load( "../examples/vision/pretrain/mnist_cnn_small.pth", map_location="cpu") model.load_state_dict(checkpoint) N = 2 n_classes = 10 image = torch.randn(N, 1, 28, 28) image = image.to(torch.float32) / 255.0 model = BoundedModule(model, torch.empty_like(image), device="cpu") eps = 0.3 norm = np.inf ptb = PerturbationLpNorm(norm=norm, eps=eps) image = BoundedTensor(image, ptb) pred = model(image) lb, ub = model.compute_bounds() assert lb.shape == ub.shape == torch.Size((2, 10)) self.result = (lb, ub) self.check()
def main(): torch.manual_seed(1234) torch.cuda.manual_seed_all(1234) random.seed(1234) np.random.seed(123) input_size = 11 ## Step 1: Initial original model as usual; note that this model has BoundedParameter as its weight parameters model_ori = SimpleNet(input_dim=input_size) ## Step 2: Prepare dataset as usual dummy_input1 = torch.randn(1, input_size) inputs = (dummy_input1, ) ## Step 3: wrap model with auto_LiRPA # The second parameter dummy_input is for constructing the trace of the computational graph. model = BoundedModule(model_ori, inputs) model.weight_perturbations = model_ori.weight_perturbations model.bias_perturbations = model_ori.bias_perturbations x = torch.randn(1, input_size) model(x) print('prediction', model_ori(x).squeeze().detach().cpu().numpy()) compute_perturbations(model, x, np.linspace(0, 0.01, 5))
shuffle=False, num_workers=4, pin_memory=True) # For illustration we only use 2 image from dataset n_classes = 100 for batch_idx, (inputs, targets) in enumerate(testloader): if True: inputs, targets = inputs.cuda(), targets.cuda() if batch_idx < 2: ## Step 3: wrap model with auto_LiRPA # The second parameter is for constructing the trace of the computational graph, and its content is not important. model = BoundedModule(model, inputs, device="cuda") ## Step 4: Compute bounds using LiRPA given a perturbation eps = 0.3 norm = np.inf ptb = PerturbationLpNorm(norm=norm, eps=eps) image = BoundedTensor(inputs, ptb) # Get model prediction as usual pred = model(image) label = torch.argmax(pred, dim=1).cpu().numpy() # Compute bounds lb, ub = model.compute_bounds() ## Step 5: Final output pred = pred.detach().cpu().numpy() lb = lb.detach().cpu().numpy()
def main(args): torch.manual_seed(args.seed) torch.cuda.manual_seed_all(args.seed) random.seed(args.seed) np.random.seed(args.seed) ## Step 1: Initial original model as usual, see model details in models/example_feedforward.py and models/example_resnet.py model_ori = models.Models[args.model]() epoch = 0 if args.load: checkpoint = torch.load(args.load) epoch, state_dict = checkpoint['epoch'], checkpoint['state_dict'] opt_state = None try: opt_state = checkpoint['optimizer'] except KeyError: print('no opt_state found') for k, v in state_dict.items(): assert torch.isnan(v).any().cpu().numpy() == 0 and torch.isinf( v).any().cpu().numpy() == 0 model_ori.load_state_dict(state_dict) logger.log('Checkpoint loaded: {}'.format(args.load)) ## Step 2: Prepare dataset as usual dummy_input = torch.randn(1, 3, 56, 56) normalize = transforms.Normalize(mean=[0.4802, 0.4481, 0.3975], std=[0.2302, 0.2265, 0.2262]) train_data = datasets.ImageFolder(args.data_dir + '/train', transform=transforms.Compose([ transforms.RandomHorizontalFlip(), transforms.RandomCrop( 56, padding_mode='edge'), transforms.ToTensor(), normalize, ])) test_data = datasets.ImageFolder( args.data_dir + '/val', transform=transforms.Compose([ # transforms.RandomResizedCrop(64, scale=(0.875, 0.875), ratio=(1., 1.)), transforms.CenterCrop(56), transforms.ToTensor(), normalize ])) train_data = torch.utils.data.DataLoader(train_data, batch_size=args.batch_size, shuffle=True, pin_memory=True, num_workers=min( multiprocessing.cpu_count(), 4)) test_data = torch.utils.data.DataLoader(test_data, batch_size=args.batch_size // 5, pin_memory=True, num_workers=min( multiprocessing.cpu_count(), 4)) train_data.mean = test_data.mean = torch.tensor([0.4802, 0.4481, 0.3975]) train_data.std = test_data.std = torch.tensor([0.2302, 0.2265, 0.2262]) ## Step 3: wrap model with auto_LiRPA # The second parameter dummy_input is for constructing the trace of the computational graph. model = BoundedModule(model_ori, dummy_input, bound_opts={'relu': args.bound_opts}, device=args.device) model_loss = BoundedModule(CrossEntropyWrapper(model_ori), (dummy_input, torch.zeros(1, dtype=torch.long)), bound_opts={ 'relu': args.bound_opts, 'loss_fusion': True }, device=args.device) model_loss = BoundDataParallel(model_loss) ## Step 4 prepare optimizer, epsilon scheduler and learning rate scheduler opt = optim.Adam(model_loss.parameters(), lr=args.lr) norm = float(args.norm) lr_scheduler = optim.lr_scheduler.MultiStepLR( opt, milestones=args.lr_decay_milestones, gamma=0.1) eps_scheduler = eval(args.scheduler_name)(args.eps, args.scheduler_opts) logger.log(str(model_ori)) if args.load: if opt_state: opt.load_state_dict(opt_state) logger.log('resume opt_state') # skip epochs if epoch > 0: epoch_length = int( (len(train_data.dataset) + train_data.batch_size - 1) / train_data.batch_size) eps_scheduler.set_epoch_length(epoch_length) eps_scheduler.train() for i in range(epoch): lr_scheduler.step() eps_scheduler.step_epoch(verbose=True) for j in range(epoch_length): eps_scheduler.step_batch() logger.log('resume from eps={:.12f}'.format(eps_scheduler.get_eps())) ## Step 5: start training if args.verify: eps_scheduler = FixedScheduler(args.eps) with torch.no_grad(): Train(model, 1, test_data, eps_scheduler, norm, False, None, 'IBP', loss_fusion=False, final_node_name=None) else: timer = 0.0 best_err = 1e10 # with torch.autograd.detect_anomaly(): for t in range(epoch + 1, args.num_epochs + 1): logger.log("Epoch {}, learning rate {}".format( t, lr_scheduler.get_last_lr())) start_time = time.time() Train(model_loss, t, train_data, eps_scheduler, norm, True, opt, args.bound_type, loss_fusion=True) lr_scheduler.step() epoch_time = time.time() - start_time timer += epoch_time logger.log('Epoch time: {:.4f}, Total time: {:.4f}'.format( epoch_time, timer)) logger.log("Evaluating...") torch.cuda.empty_cache() # remove 'model.' in state_dict state_dict_loss = model_loss.state_dict() state_dict = {} for name in state_dict_loss: assert (name.startswith('model.')) state_dict[name[6:]] = state_dict_loss[name] with torch.no_grad(): if int(eps_scheduler.params['start']) + int( eps_scheduler.params['length']) > t >= int( eps_scheduler.params['start']): m = Train(model_loss, t, test_data, eps_scheduler, norm, False, None, args.bound_type, loss_fusion=True) else: model_ori.load_state_dict(state_dict) model = BoundedModule(model_ori, dummy_input, bound_opts={'relu': args.bound_opts}, device=args.device) model = BoundDataParallel(model) m = Train(model, t, test_data, eps_scheduler, norm, False, None, 'IBP', loss_fusion=False) del model save_dict = { 'state_dict': state_dict, 'epoch': t, 'optimizer': opt.state_dict() } if t < int(eps_scheduler.params['start']): torch.save(save_dict, 'saved_models/natural_' + exp_name) elif t > int(eps_scheduler.params['start']) + int( eps_scheduler.params['length']): current_err = m.avg('Verified_Err') if current_err < best_err: best_err = current_err torch.save( save_dict, 'saved_models/' + exp_name + '_best_' + str(best_err)[:6]) else: torch.save(save_dict, 'saved_models/' + exp_name) torch.cuda.empty_cache()
test_data = torchvision.datasets.CIFAR100(root=DATAROOT, train=False, download=True, transform=transform_test) testloader = torch.utils.data.DataLoader(test_data, batch_size=64, shuffle=False, num_workers=4, pin_memory=True) n_classes = 100 for batch_idx, (image, targets) in enumerate(testloader): image, targets = image.cuda(), targets.cuda() if batch_idx < 2: print(batch_idx) model = BoundedModule(model, torch.empty_like(image), device="cuda") eps = 0.3 norm = np.inf ptb = PerturbationLpNorm(norm=norm, eps=eps) image = BoundedTensor(image, ptb) # Get model prediction as usual pred = model(image) label = torch.argmax(pred, dim=1).cpu().numpy() # Compute bounds lb, ub = model.compute_bounds() pred = pred.detach().cpu().numpy() lb = lb.detach().cpu().numpy() ub = ub.detach().cpu().numpy() for i in range(batch_idx): print("Image {} top-1 prediction {}".format(i, label[i]))
def main(args): torch.manual_seed(args.seed) torch.cuda.manual_seed_all(args.seed) random.seed(args.seed) np.random.seed(args.seed) ## Step 1: Initial original model as usual, see model details in models/example_feedforward.py and models/example_resnet.py if args.data == 'MNIST': model_ori = models.Models[args.model](in_ch=1, in_dim=28) else: model_ori = models.Models[args.model]() epoch = 0 if args.load: checkpoint = torch.load(args.load) epoch, state_dict = checkpoint['epoch'], checkpoint['state_dict'] opt_state = None try: opt_state = checkpoint['optimizer'] except KeyError: print('no opt_state found') for k, v in state_dict.items(): assert torch.isnan(v).any().cpu().numpy() == 0 and torch.isinf( v).any().cpu().numpy() == 0 model_ori.load_state_dict(state_dict) logger.log('Checkpoint loaded: {}'.format(args.load)) ## Step 2: Prepare dataset as usual if args.data == 'MNIST': dummy_input = torch.randn(1, 1, 28, 28) train_data = datasets.MNIST("./data", train=True, download=True, transform=transforms.ToTensor()) test_data = datasets.MNIST("./data", train=False, download=True, transform=transforms.ToTensor()) elif args.data == 'CIFAR': dummy_input = torch.randn(1, 3, 32, 32) normalize = transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010]) train_data = datasets.CIFAR10("./data", train=True, download=True, transform=transforms.Compose([ transforms.RandomHorizontalFlip(), transforms.RandomCrop( 32, 4, padding_mode='edge'), transforms.ToTensor(), normalize ])) test_data = datasets.CIFAR10("./data", train=False, download=True, transform=transforms.Compose( [transforms.ToTensor(), normalize])) train_data = torch.utils.data.DataLoader(train_data, batch_size=args.batch_size, shuffle=True, pin_memory=True, num_workers=min( multiprocessing.cpu_count(), 4)) test_data = torch.utils.data.DataLoader(test_data, batch_size=args.batch_size // 2, pin_memory=True, num_workers=min( multiprocessing.cpu_count(), 4)) if args.data == 'MNIST': train_data.mean = test_data.mean = torch.tensor([0.0]) train_data.std = test_data.std = torch.tensor([1.0]) elif args.data == 'CIFAR': train_data.mean = test_data.mean = torch.tensor( [0.4914, 0.4822, 0.4465]) train_data.std = test_data.std = torch.tensor([0.2023, 0.1994, 0.2010]) ## Step 3: wrap model with auto_LiRPA # The second parameter dummy_input is for constructing the trace of the computational graph. model = BoundedModule(model_ori, dummy_input, bound_opts={'relu': args.bound_opts}, device=args.device) final_name1 = model.final_name model_loss = BoundedModule(CrossEntropyWrapper(model_ori), (dummy_input, torch.zeros(1, dtype=torch.long)), bound_opts={ 'relu': args.bound_opts, 'loss_fusion': True }, device=args.device) # after CrossEntropyWrapper, the final name will change because of one additional input node in CrossEntropyWrapper final_name2 = model_loss._modules[final_name1].output_name[0] assert type(model._modules[final_name1]) == type( model_loss._modules[final_name2]) if args.no_loss_fusion: model_loss = BoundedModule(model_ori, dummy_input, bound_opts={'relu': args.bound_opts}, device=args.device) final_name2 = None model_loss = BoundDataParallel(model_loss) macs, params = profile(model_ori, (dummy_input.cuda(), )) logger.log('macs: {}, params: {}'.format(macs, params)) ## Step 4 prepare optimizer, epsilon scheduler and learning rate scheduler opt = optim.Adam(model_loss.parameters(), lr=args.lr) norm = float(args.norm) lr_scheduler = optim.lr_scheduler.MultiStepLR( opt, milestones=args.lr_decay_milestones, gamma=0.1) eps_scheduler = eval(args.scheduler_name)(args.eps, args.scheduler_opts) logger.log(str(model_ori)) # skip epochs if epoch > 0: epoch_length = int( (len(train_data.dataset) + train_data.batch_size - 1) / train_data.batch_size) eps_scheduler.set_epoch_length(epoch_length) eps_scheduler.train() for i in range(epoch): lr_scheduler.step() eps_scheduler.step_epoch(verbose=True) for j in range(epoch_length): eps_scheduler.step_batch() logger.log('resume from eps={:.12f}'.format(eps_scheduler.get_eps())) if args.load: if opt_state: opt.load_state_dict(opt_state) logger.log('resume opt_state') ## Step 5: start training if args.verify: eps_scheduler = FixedScheduler(args.eps) with torch.no_grad(): Train(model, 1, test_data, eps_scheduler, norm, False, None, 'IBP', loss_fusion=False, final_node_name=None) else: timer = 0.0 best_acc = 1e10 # with torch.autograd.detect_anomaly(): for t in range(epoch + 1, args.num_epochs + 1): logger.log("Epoch {}, learning rate {}".format( t, lr_scheduler.get_last_lr())) start_time = time.time() Train(model_loss, t, train_data, eps_scheduler, norm, True, opt, args.bound_type, loss_fusion=not args.no_loss_fusion) lr_scheduler.step() epoch_time = time.time() - start_time timer += epoch_time logger.log('Epoch time: {:.4f}, Total time: {:.4f}'.format( epoch_time, timer)) logger.log("Evaluating...") torch.cuda.empty_cache() # remove 'model.' in state_dict for CrossEntropyWrapper state_dict_loss = model_loss.state_dict() state_dict = {} if not args.no_loss_fusion: for name in state_dict_loss: assert (name.startswith('model.')) state_dict[name[6:]] = state_dict_loss[name] else: state_dict = state_dict_loss with torch.no_grad(): if t > int(eps_scheduler.params['start']) + int( eps_scheduler.params['length']): m = Train(model_loss, t, test_data, FixedScheduler(8. / 255), norm, False, None, 'IBP', loss_fusion=False, final_node_name=final_name2) else: m = Train(model_loss, t, test_data, eps_scheduler, norm, False, None, 'IBP', loss_fusion=False, final_node_name=final_name2) save_dict = { 'state_dict': state_dict, 'epoch': t, 'optimizer': opt.state_dict() } if t < int(eps_scheduler.params['start']): torch.save(save_dict, 'saved_models/natural_' + exp_name) elif t > int(eps_scheduler.params['start']) + int( eps_scheduler.params['length']): current_acc = m.avg('Verified_Err') if current_acc < best_acc: best_acc = current_acc torch.save( save_dict, 'saved_models/' + exp_name + '_best_' + str(best_acc)[:6]) else: torch.save(save_dict, 'saved_models/' + exp_name) else: torch.save(save_dict, 'saved_models/' + exp_name) torch.cuda.empty_cache()
def main(args): torch.manual_seed(args.seed) torch.cuda.manual_seed_all(args.seed) random.seed(args.seed) np.random.seed(args.seed) ## Load the model with BoundedParameter for weight perturbation. model_ori = models.Models['mlp_3layer_weight_perturb']() epoch = 0 ## Load a checkpoint, if requested. if args.load: checkpoint = torch.load(args.load) epoch, state_dict = checkpoint['epoch'], checkpoint['state_dict'] opt_state = None try: opt_state = checkpoint['optimizer'] except KeyError: print('no opt_state found') for k, v in state_dict.items(): assert torch.isnan(v).any().cpu().numpy() == 0 and torch.isinf(v).any().cpu().numpy() == 0 model_ori.load_state_dict(state_dict) logger.log('Checkpoint loaded: {}'.format(args.load)) ## Step 2: Prepare dataset as usual dummy_input = torch.randn(1, 1, 28, 28) train_data, test_data = mnist_loaders(datasets.MNIST, batch_size=args.batch_size, ratio=args.ratio) train_data.mean = test_data.mean = torch.tensor([0.0]) train_data.std = test_data.std = torch.tensor([1.0]) ## Step 3: wrap model with auto_LiRPA # The second parameter dummy_input is for constructing the trace of the computational graph. model = BoundedModule(model_ori, dummy_input, bound_opts={'relu':args.bound_opts}, device=args.device) final_name1 = model.final_name model_loss = BoundedModule(CrossEntropyWrapper(model_ori), (dummy_input, torch.zeros(1, dtype=torch.long)), bound_opts= { 'relu': args.bound_opts, 'loss_fusion': True }, device=args.device) # after CrossEntropyWrapper, the final name will change because of one more input node in CrossEntropyWrapper final_name2 = model_loss._modules[final_name1].output_name[0] assert type(model._modules[final_name1]) == type(model_loss._modules[final_name2]) if args.multigpu: model_loss = BoundDataParallel(model_loss) model_loss.ptb = model.ptb = model_ori.ptb # Perturbation on the parameters ## Step 4 prepare optimizer, epsilon scheduler and learning rate scheduler if args.opt == 'ADAM': opt = optim.Adam(model_loss.parameters(), lr=args.lr, weight_decay=0.01) elif args.opt == 'SGD': opt = optim.SGD(model_loss.parameters(), lr=args.lr, weight_decay=0.01) norm = float(args.norm) lr_scheduler = optim.lr_scheduler.MultiStepLR(opt, milestones=args.lr_decay_milestones, gamma=0.1) eps_scheduler = eval(args.scheduler_name)(args.eps, args.scheduler_opts) logger.log(str(model_ori)) # Skip epochs if we continue training from a checkpoint. if epoch > 0: epoch_length = int((len(train_data.dataset) + train_data.batch_size - 1) / train_data.batch_size) eps_scheduler.set_epoch_length(epoch_length) eps_scheduler.train() for i in range(epoch): lr_scheduler.step() eps_scheduler.step_epoch(verbose=True) for j in range(epoch_length): eps_scheduler.step_batch() logger.log('resume from eps={:.12f}'.format(eps_scheduler.get_eps())) if args.load: if opt_state: opt.load_state_dict(opt_state) logger.log('resume opt_state') ## Step 5: start training. if args.verify: eps_scheduler = FixedScheduler(args.eps) with torch.no_grad(): Train(model, 1, test_data, eps_scheduler, norm, False, None, 'CROWN-IBP', loss_fusion=False, final_node_name=None) else: timer = 0.0 best_loss = 1e10 # Main training loop for t in range(epoch + 1, args.num_epochs+1): logger.log("Epoch {}, learning rate {}".format(t, lr_scheduler.get_last_lr())) start_time = time.time() # Training one epoch Train(model_loss, t, train_data, eps_scheduler, norm, True, opt, args.bound_type, loss_fusion=True) lr_scheduler.step() epoch_time = time.time() - start_time timer += epoch_time logger.log('Epoch time: {:.4f}, Total time: {:.4f}'.format(epoch_time, timer)) logger.log("Evaluating...") torch.cuda.empty_cache() # remove 'model.' in state_dict (hack for saving models so far...) state_dict_loss = model_loss.state_dict() state_dict = {} for name in state_dict_loss: assert (name.startswith('model.')) state_dict[name[6:]] = state_dict_loss[name] # Test one epoch. with torch.no_grad(): m = Train(model_loss, t, test_data, eps_scheduler, norm, False, None, args.bound_type, loss_fusion=False, final_node_name=final_name2) # Save checkpoints. save_dict = {'state_dict': state_dict, 'epoch': t, 'optimizer': opt.state_dict()} if not os.path.exists('saved_models'): os.mkdir('saved_models') if t < int(eps_scheduler.params['start']): torch.save(save_dict, 'saved_models/natural_' + exp_name) elif t > int(eps_scheduler.params['start']) + int(eps_scheduler.params['length']): current_loss = m.avg('Loss') if current_loss < best_loss: best_loss = current_loss torch.save(save_dict, 'saved_models/' + exp_name + '_best_' + str(best_loss)[:6]) else: torch.save(save_dict, 'saved_models/' + exp_name) else: torch.save(save_dict, 'saved_models/' + exp_name) torch.cuda.empty_cache()
def __init__(self): super().__init__() self.features = BoundedModule(FeatureExtraction(), torch.empty((1, 1, 28, 28))) self.fc = nn.Linear(256, 10)
[torchvision.transforms.ToTensor(), normalize])) # For illustration we only use 1 image from dataset N = 1 n_classes = 10 image = torch.Tensor(test_data.data[:N]).reshape(N, 3, 32, 32) # Convert to float image = image.to(torch.float32) / 255.0 if device == 'cuda': image = image.cuda() ## Step 3: wrap model with auto_LiRPA # The second parameter is for constructing the trace of the computational graph, and its content is not important. # The new "patches" conv_mode provides an more efficient implementation for convolutional neural networks. model = BoundedModule(model_ori, image, bound_opts={"conv_mode": conv_mode}, device=device) ## Step 4: Compute bounds using LiRPA given a perturbation eps = 0.03 norm = np.inf ptb = PerturbationLpNorm(norm=norm, eps=eps) image = BoundedTensor(image, ptb) # Get model prediction as usual pred = model(image) # Compute bounds torch.cuda.empty_cache() print('Using {} mode to compute convolution.'.format(conv_mode)) lb, ub = model.compute_bounds(IBP=False, C=None, method='backward')
def test(): torch.manual_seed(1234) torch.cuda.manual_seed_all(1234) random.seed(1234) np.random.seed(123) model_ori = cnn_4layer_test().eval() model_ori.load_state_dict(data['model']) dummy_input = data['data'] inputs = (dummy_input, ) model = BoundedModule(model_ori, inputs) forward_ret = model(dummy_input) model_ori.eval() assert torch.isclose(model_ori(dummy_input), model(dummy_input), 1e-8).all() # Linf ptb = PerturbationLpNorm(norm=np.inf, eps=0.01) x = BoundedTensor(dummy_input, ptb) x.requires_grad_() verify_bounds(model, x, IBP=True, method=None, forward_ret=forward_ret, lb_name='l_inf_IBP_lb', ub_name='l_inf_IBP_ub') # IBP verify_bounds(model, x, IBP=True, method='backward', forward_ret=forward_ret, lb_name='l_inf_CROWN-IBP_lb', ub_name='l_inf_CROWN-IBP_ub') # CROWN-IBP verify_bounds(model, x, IBP=False, method='backward', forward_ret=forward_ret, lb_name='l_inf_CROWN_lb', ub_name='l_inf_CROWN_ub') # CROWN # L2 ptb = PerturbationLpNorm(norm=2, eps=0.01) x = BoundedTensor(dummy_input, ptb) x.requires_grad_() verify_bounds(model, x, IBP=True, method=None, forward_ret=forward_ret, lb_name='l_2_IBP_lb', ub_name='l_2_IBP_ub') # IBP verify_bounds(model, x, IBP=True, method='backward', forward_ret=forward_ret, lb_name='l_2_CROWN-IBP_lb', ub_name='l_2_CROWN-IBP_ub') # CROWN-IBP verify_bounds(model, x, IBP=False, method='backward', forward_ret=forward_ret, lb_name='l_2_CROWN_lb', ub_name='l_2_CROWN_ub') # CROWN
elif args.model == 'lstm': dummy_mask = torch.zeros(1, args.max_sent_length, device=args.device) model = LSTM(args, data_train) dev_batches = get_batches(data_dev, args.batch_size) test_batches = get_batches(data_test, args.batch_size) ptb = PerturbationSynonym(budget=args.budget) dummy_embeddings = BoundedTensor(dummy_embeddings, ptb) model_ori = model.model_from_embeddings bound_opts = {'relu': args.bound_opts_relu, 'exp': 'no-max-input'} if isinstance(model_ori, BoundedModule): model_bound = model_ori else: model_bound = BoundedModule(model_ori, (dummy_embeddings, dummy_mask), bound_opts=bound_opts, device=args.device) model.model_from_embeddings = model_bound if args.loss_fusion: bound_opts['loss_fusion'] = True model_loss = BoundedModule( CrossEntropyWrapperMultiInput(model_ori), (torch.zeros(1, dtype=torch.long), dummy_embeddings, dummy_mask), bound_opts=bound_opts, device=args.device) ptb.model = model optimizer = model.build_optimizer() if args.lr_decay < 1: lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10,
def main(args): config = load_config(args.config) logger.info('config: {}'.format(json.dumps(config))) set_seed(args.seed or config['seed']) model_ori, checkpoint, epoch, best = prepare_model(args, logger, config) logger.info('Model structure: \n {}'.format(str(model_ori))) custom_ops = {} bound_config = config['bound_params'] batch_size = (args.batch_size or config['batch_size']) test_batch_size = args.test_batch_size or batch_size dummy_input, train_data, test_data = load_data(args, config['data'], batch_size, test_batch_size, aug=not args.no_data_aug) lf = args.loss_fusion and args.bound_type == 'CROWN-IBP' bound_opts = bound_config['bound_opts'] model_ori.train() model = BoundedModule(model_ori, dummy_input, bound_opts=bound_opts, custom_ops=custom_ops, device=args.device) model_ori.to(args.device) if checkpoint is None: if args.manual_init: manual_init(args, model_ori, model, train_data) if args.kaiming_init: kaiming_init(model_ori) if lf: model_loss = BoundedModule( CrossEntropyWrapper(model_ori), (dummy_input.cuda(), torch.zeros(1, dtype=torch.long).cuda()), bound_opts=get_bound_opts_lf(bound_opts), device=args.device) params = list(model_loss.parameters()) else: model_loss = model params = list(model_ori.parameters()) logger.info('Parameter shapes: {}'.format([p.shape for p in params])) if args.multi_gpu: raise NotImplementedError('Multi-GPU is not supported yet') model = BoundDataParallel(model) model_loss = BoundDataParallel(model_loss) opt = get_optimizer(args, params, checkpoint) max_eps = args.eps or bound_config['eps'] eps_scheduler = get_eps_scheduler(args, max_eps, train_data) lr_scheduler = get_lr_scheduler(args, opt) if epoch > 0 and not args.plot: # skip epochs eps_scheduler.train() for i in range(epoch): # FIXME Can use `last_epoch` argument of lr_scheduler lr_scheduler.step() eps_scheduler.step_epoch(verbose=False) if args.verify: logger.info('Inference') meter = Train(model, model_ori, 10000, test_data, eps_scheduler, None, loss_fusion=False) logger.info(meter) else: timer = 0.0 for t in range(epoch + 1, args.num_epochs + 1): logger.info('Epoch {}, learning rate {}, dir {}'.format( t, lr_scheduler.get_last_lr(), args.dir)) start_time = time.time() if lf: Train(model_loss, model_ori, t, train_data, eps_scheduler, opt, loss_fusion=True) else: Train(model, model_ori, t, train_data, eps_scheduler, opt) update_state_dict(model_ori, model_loss) epoch_time = time.time() - start_time timer += epoch_time lr_scheduler.step() logger.info('Epoch time: {:.4f}, Total time: {:.4f}'.format( epoch_time, timer)) is_best = False if t % args.test_interval == 0: logger.info('Test without loss fusion') with torch.no_grad(): meter = Train(model, model_ori, t, test_data, eps_scheduler, None, loss_fusion=False) if eps_scheduler.get_eps() == eps_scheduler.get_max_eps(): if meter.avg('Rob_Err') < best[1]: is_best, best = True, (meter.avg('Err'), meter.avg('Rob_Err'), t) logger.info( 'Best epoch {}, error {:.4f}, robust error {:.4f}'. format(best[-1], best[0], best[1])) save(args, epoch=t, best=best, model=model_ori, opt=opt, is_best=is_best)
def main(): torch.manual_seed(1234) torch.cuda.manual_seed_all(1234) random.seed(1234) np.random.seed(123) input_size = 17 action_size = 6 ## Step 1: Initial original model as usual; note that this model has BoundedParameter as its weight parameters model_ori = RelaxedCtsPolicyForState(state_dim=input_size, action_dim=action_size) state_dict = torch.load('test_policy_net.model') if not forward_one: state_dict['log_weight'] = state_dict['log_stdev'] del state_dict['log_stdev'] # model_ori.load_state_dict(state_dict) ## Step 2: Prepare dataset as usual dummy_input1 = torch.randn(1, input_size) dummy_input2 = torch.randn(1, action_size) if forward_one: inputs = (dummy_input1, ) else: inputs = (dummy_input1, dummy_input2) model_ori(*inputs) # inputs = (dummy_input1, ) # dummy_input2 = model_ori.get_means(dummy_input1) ## Step 3: wrap model with auto_LiRPA # The second parameter dummy_input is for constructing the trace of the computational graph. model = BoundedModule(model_ori, inputs) all_states = x = torch.randn(2, input_size) means = model_ori.get_means(x).detach() if forward_one: print('prediction', model_ori(x).sum()) else: print('prediction', model_ori(x, means).sum()) action_means = means perturbations = np.arange(0.0, 0.1, 0.01) compute_perturbations(model, x, means, perturbations) if forward_one: # pred = model_ori(all_states) # pred = ((pred - means) ** 2).mean() ikl, ckl = get_kl_bound(model, all_states, action_means, 0.1) ikl, ckl = get_kl_bound(model, all_states, action_means, 0.0) print('ikl', ikl.mean().item()) print('ckl', ckl.mean().item()) pred = (0.5 * ikl + 0.5 * ckl).mean() pred.backward() print('pred', pred.item()) else: iub, cub = get_kl_bound(model, all_states, action_means, 0.1) # iub, cub = get_kl_bound(model, all_states, action_means, 0) # iub, cub = model_ori(all_states, action_means).mean() print('iub', iub.mean().item()) print('cub', cub.mean().item()) kl = (0.5 * iub + 0.5 * cub).mean() kl.backward() print('kl', kl.item()) for p in model.parameters(): if p.grad is not None: print(p.size(), p.grad.abs().sum().item()) # print(p.size(), p.grad) else: print(p.size(), p.grad)
def main(args): torch.manual_seed(args.seed) torch.cuda.manual_seed_all(args.seed) random.seed(args.seed) np.random.seed(args.seed) ## Step 1: Initial original model as usual, see model details in models/example_feedforward.py and models/example_resnet.py if args.data == 'MNIST': model_ori = models.Models[args.model](in_ch=1, in_dim=28) else: model_ori = models.Models[args.model](in_ch=3, in_dim=32) if args.load: state_dict = torch.load(args.load)['state_dict'] model_ori.load_state_dict(state_dict) ## Step 2: Prepare dataset as usual if args.data == 'MNIST': dummy_input = torch.randn(1, 1, 28, 28) train_data = datasets.MNIST("./data", train=True, download=True, transform=transforms.ToTensor()) test_data = datasets.MNIST("./data", train=False, download=True, transform=transforms.ToTensor()) elif args.data == 'CIFAR': dummy_input = torch.randn(1, 3, 32, 32) normalize = transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010]) train_data = datasets.CIFAR10("./data", train=True, download=True, transform=transforms.Compose([ transforms.RandomHorizontalFlip(), transforms.RandomCrop(32, 4), transforms.ToTensor(), normalize ])) test_data = datasets.CIFAR10("./data", train=False, download=True, transform=transforms.Compose( [transforms.ToTensor(), normalize])) train_data = torch.utils.data.DataLoader(train_data, batch_size=args.batch_size, shuffle=True, pin_memory=True, num_workers=min( multiprocessing.cpu_count(), 4)) test_data = torch.utils.data.DataLoader(test_data, batch_size=args.batch_size, pin_memory=True, num_workers=min( multiprocessing.cpu_count(), 4)) if args.data == 'MNIST': train_data.mean = test_data.mean = torch.tensor([0.0]) train_data.std = test_data.std = torch.tensor([1.0]) elif args.data == 'CIFAR': train_data.mean = test_data.mean = torch.tensor( [0.4914, 0.4822, 0.4465]) train_data.std = test_data.std = torch.tensor([0.2023, 0.1994, 0.2010]) ## Step 3: wrap model with auto_LiRPA # The second parameter dummy_input is for constructing the trace of the computational graph. model = BoundedModule(model_ori, dummy_input, bound_opts={'relu': args.bound_opts}, device=args.device) ## Step 4 prepare optimizer, epsilon scheduler and learning rate scheduler opt = optim.Adam(model.parameters(), lr=args.lr) norm = float(args.norm) lr_scheduler = optim.lr_scheduler.StepLR(opt, step_size=10, gamma=0.5) eps_scheduler = eval(args.scheduler_name)(args.eps, args.scheduler_opts) print("Model structure: \n", str(model_ori)) ## Step 5: start training if args.verify: eps_scheduler = FixedScheduler(args.eps) with torch.no_grad(): Train(model, 1, test_data, eps_scheduler, norm, False, None, args.bound_type) else: timer = 0.0 for t in range(1, args.num_epochs + 1): if eps_scheduler.reached_max_eps(): # Only decay learning rate after reaching the maximum eps lr_scheduler.step() print("Epoch {}, learning rate {}".format(t, lr_scheduler.get_lr())) start_time = time.time() Train(model, t, train_data, eps_scheduler, norm, True, opt, args.bound_type) epoch_time = time.time() - start_time timer += epoch_time print('Epoch time: {:.4f}, Total time: {:.4f}'.format( epoch_time, timer)) print("Evaluating...") with torch.no_grad(): Train(model, t, test_data, eps_scheduler, norm, False, None, args.bound_type) torch.save({ 'state_dict': model.state_dict(), 'epoch': t }, args.model)
from auto_LiRPA.perturbations import PerturbationLpNorm input_vec = torch.rand(5, 5) class mynet(nn.Module): def __init__(self): super(mynet, self).__init__() self.output = nn.sequential(nn.Linear(5, 10), nn.ReLU(), nn.Linear(10, 3)) def forward(self, input): return self.features(input) raw_model = mynet() bound_model = BoundedModule(raw_model, input_vec) num_actions = 3 batchsize = 5 label = torch.tensor([0, 2, 1, 1, 0]) bnd_state = BoundedTensor(input_vec, PerturbationLpNorm(norm=np.inf, eps=0.1)) c = torch.eye(3).type_as(input_vec)[label].unsqueeze(1) - torch.eye(3).type_as( input_vec).unsqueeze(0) I = (~(label.data.unsqueeze(1) == torch.arange(3).type_as( label.data).unsqueeze(0))) c = (c[I].view(input_vec.size(0), 2, 3)) pred = bound_model(input_vec) basic_bound, _ = bound_model.compute_bounds(IBP=False, method='backward') advance_bound, _ = bound_model.compute_bounds(C=c, IBP=False,