def testNTKGDPrediction(self, train_shape, test_shape, network, out_logits,
fn_and_kernel, momentum, learning_rate, t, loss):
key, x_test, x_train, y_train = self._get_inputs(out_logits, test_shape,
train_shape)
params, f, ntk = fn_and_kernel(key, train_shape[1:], network, out_logits)
g_dd = ntk(x_train, None, 'ntk')
g_td = ntk(x_test, x_train, 'ntk')
# Regress to an MSE loss.
loss_fn = lambda y, y_hat: 0.5 * np.mean((y - y_hat)**2)
grad_loss = jit(grad(lambda params, x: loss_fn(f(params, x), y_train)))
trace_axes = () if g_dd.ndim == 4 else (-1,)
if loss == 'mse_analytic':
if momentum is not None:
raise absltest.SkipTest(momentum)
predictor = predict.gradient_descent_mse(g_dd, y_train,
learning_rate=learning_rate,
trace_axes=trace_axes)
elif loss == 'mse':
predictor = predict.gradient_descent(loss_fn, g_dd, y_train,
learning_rate=learning_rate,
momentum=momentum,
trace_axes=trace_axes)
else:
raise NotImplementedError(loss)
predictor = jit(predictor)
fx_train_0 = f(params, x_train)
fx_test_0 = f(params, x_test)
self._test_zero_time(predictor, fx_train_0, fx_test_0, g_td, momentum)
self._test_multi_step(predictor, fx_train_0, fx_test_0, g_td, momentum)
if loss == 'mse_analytic':
self._test_inf_time(predictor, fx_train_0, fx_test_0, g_td, y_train)
if momentum is None:
opt_init, opt_update, get_params = optimizers.sgd(learning_rate)
else:
opt_init, opt_update, get_params = optimizers.momentum(learning_rate,
momentum)
opt_state = opt_init(params)
for i in range(t):
params = get_params(opt_state)
opt_state = opt_update(i, grad_loss(params, x_train), opt_state)
params = get_params(opt_state)
fx_train_nn, fx_test_nn = f(params, x_train), f(params, x_test)
fx_train_t, fx_test_t = predictor(t, fx_train_0, fx_test_0, g_td)
self.assertAllClose(fx_train_nn, fx_train_t, rtol=RTOL, atol=ATOL)
self.assertAllClose(fx_test_nn, fx_test_t, rtol=RTOL, atol=ATOL)
def testNTKGDPrediction(self, train_shape, test_shape, network, out_logits,
fn_and_kernel):
key = random.PRNGKey(0)
key, split = random.split(key)
x_train = random.normal(split, train_shape)
key, split = random.split(key)
y_train = np.array(
random.bernoulli(split, shape=(train_shape[0], out_logits)),
np.float32)
key, split = random.split(key)
x_test = random.normal(split, test_shape)
params, f, ntk = fn_and_kernel(key, train_shape[1:], network,
out_logits)
# Regress to an MSE loss.
loss = lambda y, y_hat: 0.5 * np.mean((y - y_hat)**2)
grad_loss = jit(grad(lambda params, x: loss(f(params, x), y_train)))
g_dd = ntk(x_train, None, 'ntk')
g_td = ntk(x_test, x_train, 'ntk')
predictor = predict.gradient_descent(g_dd, y_train, loss, g_td)
atol = ATOL
rtol = RTOL
step_size = 0.5
if len(train_shape) > 2:
# Hacky way to up the tolerance just for convolutions.
atol = ATOL * 2
rtol = RTOL * 2
step_size = 0.1
train_time = 100.0
steps = int(train_time / step_size)
opt_init, opt_update, get_params = optimizers.sgd(step_size)
opt_state = opt_init(params)
fx_initial_train = f(params, x_train)
fx_initial_test = f(params, x_test)
fx_pred_train, fx_pred_test = predictor(0.0, fx_initial_train,
fx_initial_test)
self.assertAllClose(fx_initial_train, fx_pred_train, True)
self.assertAllClose(fx_initial_test, fx_pred_test, True)
for i in range(steps):
params = get_params(opt_state)
opt_state = opt_update(i, grad_loss(params, x_train), opt_state)
params = get_params(opt_state)
fx_train = f(params, x_train)
fx_test = f(params, x_test)
fx_pred_train, fx_pred_test = predictor(train_time, fx_initial_train,
fx_initial_test)
fx_disp_train = np.sqrt(np.mean((fx_train - fx_initial_train)**2))
fx_disp_test = np.sqrt(np.mean((fx_test - fx_initial_test)**2))
fx_error_train = (fx_train - fx_pred_train) / fx_disp_train
fx_error_test = (fx_test - fx_pred_test) / fx_disp_test
self.assertAllClose(fx_error_train, np.zeros_like(fx_error_train),
True, rtol, atol)
self.assertAllClose(fx_error_test, np.zeros_like(fx_error_test), True,
rtol, atol)
def testPredictND(self):
n_chan = 6
key = random.PRNGKey(1)
im_shape = (5, 4, 3)
n_train = 2
n_test = 2
x_train = random.normal(key, (n_train, ) + im_shape)
y_train = random.uniform(key, (n_train, 3, 2, n_chan))
init_fn, apply_fn, _ = stax.Conv(n_chan, (3, 2), (1, 2))
_, params = init_fn(key, x_train.shape)
fx_train_0 = apply_fn(params, x_train)
for trace_axes in [(), (-1, ), (-2, ), (-3, ), (0, 1), (2, 3), (2, ),
(1, 3), (0, -1), (0, 0, -3), (0, 1, 2, 3),
(0, 1, -1, 2)]:
for ts in [None, np.arange(6).reshape((2, 3))]:
for x in [None, 'x_test']:
with self.subTest(trace_axes=trace_axes, ts=ts, x=x):
t_shape = ts.shape if ts is not None else ()
y_test_shape = t_shape + (n_test, ) + y_train.shape[1:]
y_train_shape = t_shape + y_train.shape
x = x if x is None else random.normal(
key, (n_test, ) + im_shape)
fx_test_0 = None if x is None else apply_fn(params, x)
kernel_fn = empirical.empirical_kernel_fn(
apply_fn, trace_axes=trace_axes)
# TODO(romann): investigate the SIGTERM error on CPU.
# kernel_fn = jit(kernel_fn, static_argnums=(2,))
ntk_train_train = kernel_fn(x_train, None, 'ntk',
params)
if x is not None:
ntk_test_train = kernel_fn(x, x_train, 'ntk',
params)
loss = lambda x, y: 0.5 * np.mean(x - y)**2
predict_fn_mse = predict.gradient_descent_mse(
ntk_train_train, y_train, trace_axes=trace_axes)
predict_fn_mse_ensemble = predict.gradient_descent_mse_ensemble(
kernel_fn,
x_train,
y_train,
trace_axes=trace_axes,
params=params)
if x is None:
p_train_mse = predict_fn_mse(ts, fx_train_0)
else:
p_train_mse, p_test_mse = predict_fn_mse(
ts, fx_train_0, fx_test_0, ntk_test_train)
self.assertAllClose(y_test_shape, p_test_mse.shape)
self.assertAllClose(y_train_shape, p_train_mse.shape)
p_nngp_mse_ens, p_ntk_mse_ens = predict_fn_mse_ensemble(
ts, x, ('nngp', 'ntk'), compute_cov=True)
ref_shape = y_train_shape if x is None else y_test_shape
self.assertAllClose(ref_shape,
p_ntk_mse_ens.mean.shape)
self.assertAllClose(ref_shape,
p_nngp_mse_ens.mean.shape)
if ts is not None:
predict_fn = predict.gradient_descent(
loss,
ntk_train_train,
y_train,
trace_axes=trace_axes)
if x is None:
p_train = predict_fn(ts, fx_train_0)
else:
p_train, p_test = predict_fn(
ts, fx_train_0, fx_test_0, ntk_test_train)
self.assertAllClose(y_test_shape, p_test.shape)
self.assertAllClose(y_train_shape, p_train.shape)