def main(): args = get_args() ver_logdir = args.load_model[:-3] + '_ver' if not os.path.exists(ver_logdir): os.makedirs(ver_logdir) num_train, _, test_loader, input_size, input_channel, n_class = get_loaders( args) net = get_network(device, args, input_size, input_channel, n_class) print(net) args.test_domains = [] # with torch.no_grad(): # test(device, 0, args, net, test_loader, layers=[-1, args.layer_idx]) args.test_batch = 1 num_train, _, test_loader, input_size, input_channel, n_class = get_loaders( args) latent_idx = args.layer_idx if args.latent_idx is None else args.latent_idx img_file = open(args.unverified_imgs_file, 'w') with torch.no_grad(): tot_verified_corr, tot_nat_ok, tot_attack_ok, tot_pgd_ok, tot_tests = 0, 0, 0, 0, 0 for test_idx, (inputs, targets) in enumerate(test_loader): if test_idx < args.start_idx or test_idx >= args.end_idx: continue tot_tests += 1 test_file = os.path.join(ver_logdir, '{}.p'.format(test_idx)) test_data = pickle.load(open(test_file, 'rb')) if ( not args.no_load) and os.path.isfile(test_file) else {} print('Verify test_idx =', test_idx) net.reset_bounds() inputs, targets = inputs.to(device), targets.to(device) abs_inputs = get_inputs(args.test_domain, inputs, args.test_eps, device, dtype=dtype) nat_out = net(inputs) nat_ok = targets.eq(nat_out.max(dim=1)[1]).item() tot_nat_ok += float(nat_ok) test_data['ok'] = nat_ok if not nat_ok: report(ver_logdir, tot_verified_corr, tot_nat_ok, tot_attack_ok, tot_pgd_ok, test_idx, tot_tests, test_data) continue for _ in range(args.attack_restarts): with torch.enable_grad(): pgd_loss, pgd_ok = get_adv_loss(device, args.test_eps, -1, net, None, inputs, targets, args.test_att_n_steps, args.test_att_step_size) if not pgd_ok: break if pgd_ok: test_data['pgd_ok'] = 1 tot_pgd_ok += 1 else: test_data['pgd_ok'] = 0 report(ver_logdir, tot_verified_corr, tot_nat_ok, tot_attack_ok, tot_pgd_ok, test_idx, tot_tests, test_data) continue if 'verified' in test_data and test_data['verified']: tot_verified_corr += 1 tot_attack_ok += 1 report(ver_logdir, tot_verified_corr, tot_nat_ok, tot_attack_ok, tot_pgd_ok, test_idx, tot_tests, test_data) continue if args.no_milp: report(ver_logdir, tot_verified_corr, tot_nat_ok, tot_attack_ok, tot_pgd_ok, test_idx, tot_tests, test_data) continue zono_inputs = get_inputs('zono_iter', inputs, args.test_eps, device, dtype=dtype) bounds = compute_bounds(net, device, len(net.blocks) - 1, args, zono_inputs) relu_params = reset_params(args, net, dtype) with torch.enable_grad(): learn_slopes(device, relu_params, bounds, args, len(net.blocks), net, inputs, targets, abs_inputs, None, None) bounds = compute_bounds(net, device, len(net.blocks) - 1, args, zono_inputs) for _ in range(args.attack_restarts): with torch.enable_grad(): latent_loss, latent_ok = get_adv_loss( device, args.test_eps, latent_idx, net, bounds, inputs, targets, args.test_att_n_steps, args.test_att_step_size) # print('-> ', latent_idx, latent_loss, latent_ok) if not latent_ok: break if latent_ok: tot_attack_ok += 1 zono_out = net(zono_inputs) verified, verified_corr = zono_out.verify(targets) test_data['verified'] = int(verified_corr.item()) if verified_corr: tot_verified_corr += 1 report(ver_logdir, tot_verified_corr, tot_nat_ok, tot_attack_ok, tot_pgd_ok, test_idx, tot_tests, test_data) continue loss_after = net(abs_inputs).ce_loss(targets) if args.refine_lidx is not None: bounds = compute_bounds(net, device, len(net.blocks) - 1, args, abs_inputs) for lidx in range(0, args.layer_idx + 2): net.blocks[lidx].bounds = bounds[lidx] print('loss before refine: ', net(abs_inputs).ce_loss(targets)) refine_dim = bounds[args.refine_lidx + 1][0].shape[2] pbar = tqdm(total=refine_dim * refine_dim, dynamic_ncols=True) for refine_i in range(refine_dim): for refine_j in range(refine_dim): refine(args, bounds, net, refine_i, refine_j, abs_inputs, input_size) pbar.update(1) pbar.close() loss_after = net(abs_inputs).ce_loss(targets) print('loss after refine: ', loss_after) if loss_after < args.loss_threshold: if args.refine_opt is not None: with torch.enable_grad(): learn_bounds(net, bounds, relu_params, zono_inputs, args.refine_opt) if verify_test(args, net, inputs, targets, abs_inputs, bounds, test_data, test_idx): tot_verified_corr += 1 test_data['verified'] = True report(ver_logdir, tot_verified_corr, tot_nat_ok, tot_attack_ok, tot_pgd_ok, test_idx, tot_tests, test_data) img_file.close()
def verify_test(args, net, inputs, targets, abs_inputs, bounds, test_data, test_idx): ok = True n_layers = len(net.blocks) for adv_idx in range(10): if targets[0] == adv_idx: continue if ('verified', adv_idx) in test_data and test_data[('verified', adv_idx)]: print('label already verified: ', adv_idx) continue relu_params = reset_params(args, net, dtype) bin_neurons = None with torch.enable_grad(): verified, relu_priority = learn_slopes(device, relu_params, bounds, args, n_layers, net, inputs, targets, abs_inputs, targets[0].item(), adv_idx) if args.tot_binary is not None: bin_neurons = {} all_neurons = [] for layer_idx, neurons in relu_priority.items(): for neuron_idx, neuron_priority in enumerate(neurons): all_neurons += [(layer_idx, neuron_idx, neuron_priority)] # print(layer_idx, neuron_idx, neuron_priority) all_neurons.sort(key=lambda x: -x[2]) bin_neurons = {(layer_idx, neuron_idx): True for layer_idx, neuron_idx, neuron_priority in all_neurons[:args.tot_binary]} if verified: print('adv_idx=%d verified without MILP' % adv_idx) test_data[('verified', adv_idx)] = True continue model = Model("milp") model.setParam('OutputFlag', args.debug) model.setParam('TimeLimit', args.milp_timeout) abs_curr = net.forward_until(args.layer_idx, abs_inputs) abs_flat = abs_curr.view((1, -1)) n_inputs = abs_flat.head.size()[1] if abs_flat.errors is not None: n_errors = abs_flat.errors.size()[0] errors = [ model.addVar(-1.0, 1.0, vtype=GRB.CONTINUOUS, name='error_{}'.format(j)) for j in range(n_errors) ] lb, ub = abs_flat.concretize() lb, ub = lb.detach().cpu().numpy(), ub.detach().cpu().numpy() neurons = {} neurons[args.layer_idx] = [] for j in range(n_inputs): neurons[args.layer_idx] += [ model.addVar(lb=lb[0, j], ub=ub[0, j], vtype=GRB.CONTINUOUS, name='input_{}'.format(j)) ] expr = LinExpr() expr += abs_flat.head[0, j].item() expr += LinExpr( abs_flat.errors[:, 0, j].detach().cpu().numpy().tolist(), errors) model.addConstr(expr, GRB.EQUAL, neurons[args.layer_idx][j]) n_outs = n_inputs for lidx in range(args.layer_idx + 1, n_layers): pr_lb, pr_ub = lb, ub abs_curr = net.blocks[lidx](abs_curr) abs_flat = abs_curr.view((1, -1)) neurons[lidx] = [] lb, ub = abs_flat.concretize() lb, ub = lb.detach().cpu().numpy(), ub.detach().cpu().numpy() if isinstance(net.blocks[lidx], Linear): weight, bias = net.blocks[lidx].linear.weight, net.blocks[ lidx].linear.bias n_outs = weight.size()[0] for out_idx in range(n_outs): nvar = model.addVar(vtype=GRB.CONTINUOUS, lb=lb[0, out_idx], ub=ub[0, out_idx], name='n_{}_{}'.format(lidx, out_idx)) neurons[lidx].append(nvar) tmp = LinExpr() tmp += -neurons[lidx][out_idx] tmp += bias[out_idx].item() tmp += LinExpr(weight[out_idx].detach().cpu().numpy(), neurons[lidx - 1]) model.addConstr(tmp, GRB.EQUAL, 0) elif isinstance(net.blocks[lidx], ReLU): if net.blocks[lidx].bounds is not None: pr_lb = np.maximum( pr_lb, net.blocks[lidx].bounds[0].view( (1, -1)).cpu().numpy()) pr_ub = np.minimum( pr_ub, net.blocks[lidx].bounds[1].view( (1, -1)).cpu().numpy()) unstable, n_binary = handle_relu(args.max_binary, lidx, relu_priority[lidx], model, neurons, n_outs, pr_lb, pr_ub, bin_neurons=bin_neurons) print('Unstable ReLU: ', unstable, ' binary: ', n_binary) elif isinstance(net.blocks[lidx], Flatten): # print('flatten') neurons[lidx] = neurons[lidx - 1] elif isinstance(net.blocks[lidx], Conv2d): img_dim = abs_curr.head.size()[2] weight, bias = net.blocks[lidx].conv.weight.cpu().numpy( ), net.blocks[lidx].conv.bias.cpu().numpy() kernel_size, stride, pad = net.blocks[ lidx].kernel_size, net.blocks[lidx].stride, net.blocks[ lidx].padding out_channels, in_channels = weight.shape[0], weight.shape[1] neurons[lidx] = [] for out_ch in range(out_channels): for x in range(0, img_dim): for y in range(0, img_dim): new_idx = out_ch * (img_dim**2) + x * img_dim + y nvar = model.addVar(vtype=GRB.CONTINUOUS, lb=lb[0, new_idx], ub=ub[0, new_idx], name='n_{}_{}'.format( lidx, new_idx)) neurons[lidx].append(nvar) expr = LinExpr() expr += bias[out_ch] for kx in range(0, kernel_size): for ky in range(0, kernel_size): new_x = -pad + x * stride + kx new_y = -pad + y * stride + ky if new_x < 0 or new_y < 0 or new_x >= net.blocks[ lidx].bounds.size( 2) or new_y >= net.blocks[ lidx].bounds.size(3): continue for in_ch in range(in_channels): old_idx = in_ch * (net.blocks[lidx].bounds.size(2) * net.blocks[lidx].bounds.size(3)) \ + new_x * (net.blocks[lidx].bounds.size(3)) + new_y expr += neurons[ lidx - 1][old_idx] * weight[out_ch, in_ch, kx, ky] model.addConstr(expr, GRB.EQUAL, nvar) else: print('unknown layer type: ', net.blocks[lidx]) assert False model.setObjective( neurons[n_layers - 1][targets[0].item()] - neurons[n_layers - 1][adv_idx], GRB.MINIMIZE) model.update() model.optimize(callback) print('MILP: ', targets[0].item(), adv_idx, model.objVal, model.objBound, model.RunTime) test_data[adv_idx] = { 'milp_timeout': args.milp_timeout, 'max_binary': args.max_binary, 'obj_val': model.objVal, 'obj_bound': model.objBound, 'runtime': model.RunTime, } if model.objBound < 0: ok = False if args.fail_break: break else: test_data[('verified', adv_idx)] = True return ok