def test_abs(self, same_inputs):
key = random.PRNGKey(1)
X0_1 = random.normal(key, (5, 7))
fc = stax.Dense(10, 1, 0)
X0_2 = None if same_inputs else random.normal(key, (9, 7))
# Test that Abs == ABRelu(-1, 1)
init_fn, apply_leaky_relu, kernel_fn_abs = stax.serial(fc, stax.Abs())
_, apply_ab_relu, kernel_fn_ab_relu = stax.serial(fc, stax.ABRelu(-1, 1))
params = init_fn(key, input_shape=(-1, 7))
X1_1_abs = apply_leaky_relu(params, X0_1)
X1_1_ab_relu = apply_ab_relu(params, X0_1)
self.assertAllClose(X1_1_abs, X1_1_ab_relu, True)
kernels_abs = kernel_fn_abs(X0_1, X0_2, ('nngp', 'ntk'))
kernels_ab_relu = kernel_fn_ab_relu(X0_1, X0_2, ('nngp', 'ntk'))
self.assertAllClose(kernels_abs, kernels_ab_relu, True)
def test_abs(self, same_inputs, do_stabilize):
key = random.PRNGKey(1)
X0_1 = random.normal(key, (3, 2))
fc = stax.Dense(5, 1, 0)
X0_2 = None if same_inputs else random.normal(key, (4, 2))
# Test that Abs == ABRelu(-1, 1)
init_fn, apply_leaky_relu, kernel_fn_abs = stax.serial(
fc, stax.Abs(do_stabilize=do_stabilize))
_, apply_ab_relu, kernel_fn_ab_relu = stax.serial(fc, stax.ABRelu(-1, 1))
_, params = init_fn(key, input_shape=X0_1.shape)
X1_1_abs = apply_leaky_relu(params, X0_1)
X1_1_ab_relu = apply_ab_relu(params, X0_1)
self.assertAllClose(X1_1_abs, X1_1_ab_relu)
kernels_abs = kernel_fn_abs(X0_1, X0_2, ('nngp', 'ntk'))
kernels_ab_relu = kernel_fn_ab_relu(X0_1, X0_2, ('nngp', 'ntk'))
self.assertAllClose(kernels_abs, kernels_ab_relu)
def _get_phi(cls, i):
return {0: stax.Relu(), 1: stax.Erf(), 2: stax.Abs()}[i % 3]
def test_mask_conv(self, same_inputs, get, mask_axis, mask_constant,
concat, proj, p, n, transpose):
if isinstance(concat, int) and concat > n:
raise absltest.SkipTest('Concatenation axis out of bounds.')
test_utils.skip_test(self)
if default_backend() == 'gpu' and n > 3:
raise absltest.SkipTest('>=4D-CNN is not supported on GPUs.')
width = 256
n_samples = 256
tol = 0.03
key = random.PRNGKey(1)
spatial_shape = ((1, 2, 3, 2, 1) if transpose else (15, 8, 9))[:n]
filter_shape = ((2, 3, 1, 2, 1) if transpose else (7, 2, 3))[:n]
strides = (2, 1, 3, 2, 3)[:n]
spatial_spec = 'HWDZX'[:n]
dimension_numbers = ('N' + spatial_spec + 'C', 'OI' + spatial_spec,
'N' + spatial_spec + 'C')
x1 = np.cos(random.normal(key, (2, ) + spatial_shape + (2, )))
x1 = test_utils.mask(x1, mask_constant, mask_axis, key, p)
if same_inputs:
x2 = None
else:
x2 = np.cos(random.normal(key, (4, ) + spatial_shape + (2, )))
x2 = test_utils.mask(x2, mask_constant, mask_axis, key, p)
def get_attn():
return stax.GlobalSelfAttention(
n_chan_out=width,
n_chan_key=width,
n_chan_val=int(np.round(float(width) / int(np.sqrt(width)))),
n_heads=int(np.sqrt(width)),
) if proj == 'avg' else stax.Identity()
conv = stax.ConvTranspose if transpose else stax.Conv
nn = stax.serial(
stax.FanOut(3),
stax.parallel(
stax.serial(
conv(dimension_numbers=dimension_numbers,
out_chan=width,
strides=strides,
filter_shape=filter_shape,
padding='CIRCULAR',
W_std=1.5,
b_std=0.2),
stax.LayerNorm(axis=(1, -1)),
stax.Abs(),
stax.DotGeneral(rhs=0.9),
conv(dimension_numbers=dimension_numbers,
out_chan=width,
strides=strides,
filter_shape=filter_shape,
padding='VALID',
W_std=1.2,
b_std=0.1),
),
stax.serial(
conv(dimension_numbers=dimension_numbers,
out_chan=width,
strides=strides,
filter_shape=filter_shape,
padding='SAME',
W_std=0.1,
b_std=0.3),
stax.Relu(),
stax.Dropout(0.7),
conv(dimension_numbers=dimension_numbers,
out_chan=width,
strides=strides,
filter_shape=filter_shape,
padding='VALID',
W_std=0.9,
b_std=1.),
),
stax.serial(
get_attn(),
conv(dimension_numbers=dimension_numbers,
out_chan=width,
strides=strides,
filter_shape=filter_shape,
padding='CIRCULAR',
W_std=1.,
b_std=0.1),
stax.Erf(),
stax.Dropout(0.2),
stax.DotGeneral(rhs=0.7),
conv(dimension_numbers=dimension_numbers,
out_chan=width,
strides=strides,
filter_shape=filter_shape,
padding='VALID',
W_std=1.,
b_std=0.1),
)),
(stax.FanInSum() if concat is None else stax.FanInConcat(concat)),
get_attn(),
{
'avg': stax.GlobalAvgPool(),
'sum': stax.GlobalSumPool(),
'flatten': stax.Flatten(),
}[proj],
)
if get == 'nngp':
init_fn, apply_fn, kernel_fn = stax.serial(
nn, stax.Dense(width, 1., 0.))
elif get == 'ntk':
init_fn, apply_fn, kernel_fn = stax.serial(nn,
stax.Dense(1, 1., 0.))
else:
raise ValueError(get)
kernel_fn_mc = nt.monte_carlo_kernel_fn(
init_fn,
apply_fn,
key,
n_samples,
device_count=0 if concat in (0, -n) else -1,
implementation=_DEFAULT_TESTING_NTK_IMPLEMENTATION,
vmap_axes=None if concat in (0, -n) else 0,
)
kernel_fn = jit(kernel_fn, static_argnames='get')
exact = kernel_fn(x1, x2, get, mask_constant=mask_constant)
empirical = kernel_fn_mc(x1, x2, get=get, mask_constant=mask_constant)
test_utils.assert_close_matrices(self, empirical, exact, tol)
def test_mask_fc(self, same_inputs, get, concat, p, mask_axis,
mask_constant):
width = 512
n_samples = 128
tol = 0.04
key = random.PRNGKey(1)
x1 = random.normal(key, (4, 6, 5, 7))
x1 = test_utils.mask(x1, mask_constant, mask_axis, key, p)
if same_inputs:
x2 = None
else:
x2 = random.normal(key, (2, 6, 5, 7))
x2 = test_utils.mask(x2, mask_constant, mask_axis, key, p)
nn = stax.serial(
stax.Flatten(), stax.FanOut(3),
stax.parallel(
stax.serial(
stax.Dense(width, 1., 0.1),
stax.Abs(),
stax.DotGeneral(lhs=-0.2),
stax.Dense(width, 1.5, 0.01),
),
stax.serial(
stax.Dense(width, 1.1, 0.1),
stax.DotGeneral(rhs=0.7),
stax.Erf(),
stax.Dense(width if concat != 1 else 512, 1.5, 0.1),
),
stax.serial(
stax.DotGeneral(rhs=0.5),
stax.Dense(width, 1.2),
stax.ABRelu(-0.2, 0.4),
stax.Dense(width if concat != 1 else 1024, 1.3, 0.2),
)),
(stax.FanInSum() if concat is None else stax.FanInConcat(concat)),
stax.Dense(width, 2., 0.01), stax.Relu())
if get == 'nngp':
init_fn, apply_fn, kernel_fn = stax.serial(
nn, stax.Dense(width, 1., 0.1))
elif get == 'ntk':
init_fn, apply_fn, kernel_fn = stax.serial(nn,
stax.Dense(1, 1., 0.1))
else:
raise ValueError(get)
kernel_fn_mc = nt.monte_carlo_kernel_fn(
init_fn,
apply_fn,
key,
n_samples,
device_count=0 if concat in (0, -2) else -1,
implementation=_DEFAULT_TESTING_NTK_IMPLEMENTATION,
vmap_axes=None if concat in (0, -2) else 0,
)
kernel_fn = jit(kernel_fn, static_argnames='get')
exact = kernel_fn(x1, x2, get, mask_constant=mask_constant)
empirical = kernel_fn_mc(x1, x2, get=get, mask_constant=mask_constant)
test_utils.assert_close_matrices(self, empirical, exact, tol)