def test_exp_fastlin(self): def exp_model(inp): return jnp.exp(inp) exp_inp_shape = (4, 7) lb, ub = test_utils.sample_bounds(jax.random.PRNGKey(0), exp_inp_shape, minval=-10., maxval=10.) input_bounds = jax_verify.IntervalBound(lb, ub) output_bounds = jax_verify.forward_fastlin_bound_propagation( exp_model, input_bounds) uniform_inps = test_utils.sample_bounded_points( jax.random.PRNGKey(1), (lb, ub), 100) uniform_outs = jax.vmap(exp_model)(uniform_inps) empirical_min = uniform_outs.min(axis=0) empirical_max = uniform_outs.max(axis=0) self.assertGreaterEqual((output_bounds.upper - empirical_max).min(), 0., 'Invalid upper bound for Exponential. The gap ' 'between upper bound and empirical max is < 0') self.assertGreaterEqual( (empirical_min - output_bounds.lower).min(), 0., 'Invalid lower bound for Exponential. The gap' 'between emp. min and lower bound is negative.')
def test_fastlin(self, model_cls): @hk.transform_with_state def model_pred(inputs, is_training, test_local_stats=False): model = model_cls() return model(inputs, is_training, test_local_stats) inps = jnp.zeros((4, 28, 28, 1), dtype=jnp.float32) params, state = model_pred.init(jax.random.PRNGKey(42), inps, is_training=True) def logits_fun(inputs): return model_pred.apply(params, state, None, inputs, False, test_local_stats=False)[0] input_bounds = jax_verify.IntervalBound(inps - 1.0, inps + 1.0) jax_verify.forward_fastlin_bound_propagation(logits_fun, input_bounds)
def test_relu_fixed_fastlin(self): def relu_model(inp): return jax.nn.relu(inp) z = jnp.array([[-2., 3.]]) input_bounds = jax_verify.IntervalBound(z - 1., z + 1.) output_bounds = jax_verify.forward_fastlin_bound_propagation( relu_model, input_bounds) self.assertArrayAlmostEqual(jnp.array([[0., 2.]]), output_bounds.lower) self.assertArrayAlmostEqual(jnp.array([[0., 4.]]), output_bounds.upper)
def main(unused_args): # Load the parameters of an existing model. model_pred, params = load_model(FLAGS.model) logits_fn = functools.partial(model_pred, params) # Load some test samples with jax_verify.open_file('mnist/x_test_first100.npy', 'rb') as f: inputs = np.load(f) # Compute boundprop bounds eps = 0.1 lower_bound = jnp.minimum(jnp.maximum(inputs[:2, ...] - eps, 0.0), 1.0) upper_bound = jnp.minimum(jnp.maximum(inputs[:2, ...] + eps, 0.0), 1.0) init_bound = jax_verify.IntervalBound(lower_bound, upper_bound) if FLAGS.boundprop_method == 'forwardfastlin': final_bound = jax_verify.forward_fastlin_bound_propagation(logits_fn, init_bound) boundprop_transform = forward_linear_bounds.forward_fastlin_transform elif FLAGS.boundprop_method == 'ibp': final_bound = jax_verify.interval_bound_propagation(logits_fn, init_bound) boundprop_transform = jax_verify.ibp_transform else: raise NotImplementedError('Only ibp/fastlin boundprop are' 'currently supported') dummy_output = model_pred(params, inputs) # Run LP solver objective = jnp.where(jnp.arange(dummy_output[0, ...].size) == 0, jnp.ones_like(dummy_output[0, ...]), jnp.zeros_like(dummy_output[0, ...])) objective_bias = 0. value, _, status = jax_verify.solve_planet_relaxation( logits_fn, init_bound, boundprop_transform, objective, objective_bias, index=0) logging.info('Relaxation LB is : %f, Status is %s', value, status) value, _, status = jax_verify.solve_planet_relaxation( logits_fn, init_bound, boundprop_transform, -objective, objective_bias, index=0) logging.info('Relaxation UB is : %f, Status is %s', -value, status) logging.info('Boundprop LB is : %f', final_bound.lower[0, 0]) logging.info('Boundprop UB is : %f', final_bound.upper[0, 0])
def _compute_jv_bounds( input_bound: sdp_utils.IntBound, params: ModelParams, method: str, ) -> List[sdp_utils.IntBound]: """Compute bounds with jax_verify.""" jv_input_bound = jax_verify.IntervalBound(input_bound.lb, input_bound.ub) # create a function that takes as arguments the input and all parameters # that have bounds (as specified in param_bounds) and returns all # activations all_act_fun = _make_all_act_fn(params) # use jax_verify to perform (bilinear) interval bound propagation jv_param_bounds = [(p.w_bound, p.b_bound) for p in params if p.has_bounds] if method == 'ibp': _, jv_bounds = jax_verify.interval_bound_propagation( all_act_fun, jv_input_bound, *jv_param_bounds) elif method == 'fastlin': _, jv_bounds = jax_verify.forward_fastlin_bound_propagation( all_act_fun, jv_input_bound, *jv_param_bounds) elif method == 'ibpfastlin': _, jv_bounds = jax_verify.ibpforwardfastlin_bound_propagation( all_act_fun, jv_input_bound, *jv_param_bounds) elif method == 'crown': _, jv_bounds = jax_verify.backward_crown_bound_propagation( all_act_fun, jv_input_bound, *jv_param_bounds) elif method == 'nonconvex': _, jv_bounds = jax_verify.nonconvex_constopt_bound_propagation( all_act_fun, jv_input_bound, *jv_param_bounds) else: raise ValueError('Unsupported method.') # re-format bounds with internal convention bounds = [] for intermediate_bound in jv_bounds: bounds.append( sdp_utils.IntBound(lb_pre=intermediate_bound.lower, ub_pre=intermediate_bound.upper, lb=jnp.maximum(intermediate_bound.lower, 0), ub=jnp.maximum(intermediate_bound.upper, 0))) return bounds
def test_nobatch_batch_inputs(self): batch_shape = (3, 2) unbatch_shape = (2, 4) def bilinear_model(inp_1, inp_2): return jnp.einsum('bh,hH->bH', inp_1, inp_2) lb_1, ub_1 = test_utils.sample_bounds(jax.random.PRNGKey(0), batch_shape, minval=-10, maxval=10.) lb_2, ub_2 = test_utils.sample_bounds(jax.random.PRNGKey(1), unbatch_shape, minval=-10, maxval=10.) bound_1 = jax_verify.IntervalBound(lb_1, ub_1) bound_2 = jax_verify.IntervalBound(lb_2, ub_2) output_bounds = jax_verify.forward_fastlin_bound_propagation( bilinear_model, bound_1, bound_2) uniform_1 = test_utils.sample_bounded_points(jax.random.PRNGKey(2), (lb_1, ub_1), 100) uniform_2 = test_utils.sample_bounded_points(jax.random.PRNGKey(3), (lb_2, ub_2), 100) uniform_outs = jax.vmap(bilinear_model)(uniform_1, uniform_2) empirical_min = uniform_outs.min(axis=0) empirical_max = uniform_outs.max(axis=0) self.assertGreaterEqual( (output_bounds.upper - empirical_max).min(), 0., 'Invalid upper bound for mix of batched/unbatched' 'input bounds.') self.assertGreaterEqual( (empirical_min - output_bounds.lower).min(), 0., 'Invalid lower bound for mix of batched/unbatched' 'input bounds.')
def test_multiply_fastlin(self): def multiply_model(lhs, rhs): return lhs * rhs mul_inp_shape = (4, 7) lhs_lb, lhs_ub = test_utils.sample_bounds(jax.random.PRNGKey(0), mul_inp_shape, minval=-10., maxval=10.) rhs_lb, rhs_ub = test_utils.sample_bounds(jax.random.PRNGKey(1), mul_inp_shape, minval=-10., maxval=10.) lhs_bounds = jax_verify.IntervalBound(lhs_lb, lhs_ub) rhs_bounds = jax_verify.IntervalBound(rhs_lb, rhs_ub) output_bounds = jax_verify.forward_fastlin_bound_propagation( multiply_model, lhs_bounds, rhs_bounds) uniform_lhs_inps = test_utils.sample_bounded_points( jax.random.PRNGKey(2), (lhs_lb, lhs_ub), 100) uniform_rhs_inps = test_utils.sample_bounded_points( jax.random.PRNGKey(3), (rhs_lb, rhs_ub), 100) uniform_outs = jax.vmap(multiply_model)(uniform_lhs_inps, uniform_rhs_inps) empirical_min = uniform_outs.min(axis=0) empirical_max = uniform_outs.max(axis=0) self.assertGreaterEqual( (output_bounds.upper - empirical_max).min(), 0., 'Invalid upper bound for Multiply. The gap ' 'between upper bound and empirical max is negative') self.assertGreaterEqual( (empirical_min - output_bounds.lower).min(), 0., 'Invalid lower bound for Multiply. The gap' 'between emp. min and lower bound is negative.')