def f():
stateless.stateless_random_uniform(shape=shape,
seed=[1, 2],
minval=array_ops.zeros(
shape, 'int32'),
maxval=100,
dtype='int32')
def f2():
stateless.stateless_random_uniform(
shape=shape,
seed=[1, 2],
minval=0,
maxval=array_ops.ones(shape, 'int32') * 100,
dtype='int32')
def stateless_split(seed, num=2):
"""Splits an RNG seed into `num` new seeds by adding a leading axis.
Example:
>>> seed = [1, 2]
>>> new_seeds = tf.random.experimental.stateless_split(seed, num=3)
>>> print(new_seeds)
tf.Tensor(
[[1105988140 1738052849]
[-335576002 370444179]
[ 10670227 -246211131]], shape=(3, 2), dtype=int32)
>>> tf.random.stateless_normal(shape=[3], seed=new_seeds[0, :])
<tf.Tensor: shape=(3,), dtype=float32, numpy=array([-0.59835213, -0.9578608 ,
0.9002807 ], dtype=float32)>
Args:
seed: an RNG seed (a tensor with shape [2] and dtype `int32` or
`int64`). (When using XLA, only `int32` is allowed.)
num: optional, a positive integer or scalar tensor indicating the number of
seeds to produce (default 2).
Returns:
A tensor with shape [num, 2] representing `num` new seeds. It will have the
same dtype as `seed` (if `seed` doesn't have an explict dtype, the dtype
will be determined by `tf.convert_to_tensor`).
"""
seed = ops.convert_to_tensor(seed)
return stateless_random_ops.stateless_random_uniform(shape=[num, 2],
seed=seed,
dtype=seed.dtype,
minval=None,
maxval=None)
def builder_fn():
shape = (10, 1000, 1000)
seed_var = variables.Variable((312, 456),
dtype=dtypes.int32,
name='input')
random_t = stateless.stateless_random_uniform(
shape, seed=seed_var, dtype=dtype)
return '%s.shape%s' % (name, shape), [random_t]
def _random_uniform(self, *args, **kwargs):
if self._use_stateless:
c_seed = self._stateless_seed_offset + kwargs['seed']
kwargs['seed'] = math_ops.cast(
array_ops.stack([c_seed, self._global_step]), dtypes.int32)
return stateless_random_ops.stateless_random_uniform(
*args, **kwargs)
else:
return random_ops.random_uniform(*args, **kwargs)
def testRandomUniformIsInRange(self):
with self.cached_session() as sess, self.test_scope():
for dtype in self._random_types():
seed_t = array_ops.placeholder(dtypes.int32, shape=[2])
x = stateless.stateless_random_uniform(
shape=[1000], seed=seed_t, dtype=dtype)
y = sess.run(x, {seed_t: [0x12345678, 0xabcdef12]})
self.assertTrue(np.all(y >= 0))
self.assertTrue(np.all(y < 1))
def testRandomUniformIsInRange(self):
with self.session() as sess, self.test_scope():
for dtype in self._random_types(include_int=True):
maxval = 1
if dtype.is_integer:
maxval = 100
seed_t = array_ops.placeholder(dtypes.int32, shape=[2])
x = stateless.stateless_random_uniform(
shape=[1000], seed=seed_t, maxval=maxval, dtype=dtype)
y = sess.run(x, {seed_t: [0x12345678, 0xabcdef1]})
self.assertTrue(np.all(y >= 0))
self.assertTrue(np.all(y < maxval))
def testNDimensional(self):
with self.test_scope():
shape = [77, 10, 3, 5, 7]
superdiag = random.stateless_random_uniform(
shape=shape[:-1], dtype=dtypes.float32, seed=[5, 10])
maindiag = random.stateless_random_uniform(
shape=shape[:-1], dtype=dtypes.float32, seed=[5, 11])
subdiag = random.stateless_random_uniform(
shape=shape[:-1], dtype=dtypes.float32, seed=[5, 12])
rhs = random.stateless_random_uniform(
shape=shape, dtype=dtypes.float32, seed=[5, 13])
expected = self._tridiagonal_matmul((superdiag, maindiag, subdiag),
rhs,
diagonals_format='sequence')
real = self._jit_tridiagonal_matmul((superdiag, maindiag, subdiag),
rhs,
diagonals_format='sequence')
self.assertAllClose(expected, real)
def main(unused_argv):
# Build data pipelines.
print('Loading data.')
x_train, y_train, x_test, y_test = \
datasets.get_dataset('mnist', FLAGS.train_size, FLAGS.test_size)
# Build the network
init_fn, apply_fn, _ = stax.serial(stax.Dense(512, 1., 0.05), stax.Erf(),
stax.Dense(10, 1., 0.05))
key = random.stateless_random_uniform(shape=[2],
seed=[0, 0],
minval=None,
maxval=None,
dtype=np.int32)
_, params = init_fn(key, (1, 784))
# Create and initialize an optimizer.
opt_init, opt_apply, get_params = optimizers.sgd(FLAGS.learning_rate)
state = opt_init(params)
# Create an mse loss function and a gradient function.
loss = lambda fx, y_hat: 0.5 * np.mean((fx - y_hat)**2)
grad_loss = jit(grad(lambda params, x, y: loss(apply_fn(params, x), y)))
# Create an MSE predictor to solve the NTK equation in function space.
ntk = nt.batch(nt.empirical_ntk_fn(apply_fn), batch_size=4, device_count=0)
g_dd = ntk(x_train, None, params)
g_td = ntk(x_test, x_train, params)
predictor = nt.predict.gradient_descent_mse(g_dd, y_train)
# Get initial values of the network in function space.
fx_train = apply_fn(params, x_train)
fx_test = apply_fn(params, x_test)
# Train the network.
train_steps = int(FLAGS.train_time // FLAGS.learning_rate)
print('Training for {} steps'.format(train_steps))
for i in range(train_steps):
params = get_params(state)
state = opt_apply(i, grad_loss(params, x_train, y_train), state)
# Get predictions from analytic computation.
print('Computing analytic prediction.')
fx_train, fx_test = predictor(FLAGS.train_time, fx_train, fx_test, g_td)
# Print out summary data comparing the linear / nonlinear model.
util.print_summary('train', y_train, apply_fn(params, x_train), fx_train,
loss)
util.print_summary('test', y_test, apply_fn(params, x_test), fx_test, loss)
def testDistributionOfStatelessRandomUniform(self):
"""Use Pearson's Chi-squared test to test for uniformity."""
with self.cached_session() as sess, self.test_scope():
for dtype in self._random_types():
seed_t = array_ops.placeholder(dtypes.int32, shape=[2])
n = 1000
x = stateless.stateless_random_uniform(
shape=[n], seed=seed_t, dtype=dtype)
y = sess.run(x, {seed_t: [565656, 121212]})
# Tests that the values are distributed amongst 10 bins with equal
# probability. 16.92 is the Chi^2 value for 9 degrees of freedom with
# p=0.05. This test is probabilistic and would be flaky if the random
# seed were not fixed.
self.assertTrue(self._chi_squared(y, 10) < 16.92)
def stateless_random_crop(value, size, seed, name=None):
"""Randomly crops a tensor to a given size in a deterministic manner.
Slices a shape `size` portion out of `value` at a uniformly chosen offset.
Requires `value.shape >= size`.
If a dimension should not be cropped, pass the full size of that dimension.
For example, RGB images can be cropped with
`size = [crop_height, crop_width, 3]`.
Guarantees the same results given the same `seed` independent of how many
times the function is called, and independent of global seed settings (e.g.
`tf.random.set_seed`).
Usage Example:
>>> image = [[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]]
>>> seed = (1, 2)
>>> tf.image.stateless_random_crop(value=image, size=(1, 2, 3), seed=seed)
<tf.Tensor: shape=(1, 2, 3), dtype=int32, numpy=
array([[[1, 2, 3],
[4, 5, 6]]], dtype=int32)>
Args:
value: Input tensor to crop.
size: 1-D tensor with size the rank of `value`.
seed: A shape [2] Tensor, the seed to the random number generator. Must have
dtype `int32` or `int64`. (When using XLA, only `int32` is allowed.)
name: A name for this operation (optional).
Returns:
A cropped tensor of the same rank as `value` and shape `size`.
"""
with ops.name_scope(name, "random_crop", [value, size]) as name:
value = ops.convert_to_tensor(value, name="value")
size = ops.convert_to_tensor(size, dtype=dtypes.int32, name="size")
shape = array_ops.shape(value)
check = control_flow_ops.Assert(
math_ops.reduce_all(shape >= size),
["Need value.shape >= size, got ", shape, size],
summarize=1000)
shape = control_flow_ops.with_dependencies([check], shape)
limit = shape - size + 1
offset = stateless_random_ops.stateless_random_uniform(
array_ops.shape(shape),
dtype=size.dtype,
maxval=size.dtype.max,
seed=seed) % limit
return array_ops.slice(value, offset, size, name=name)
def random_cropped_inputs():
"""Cropped inputs with stateless random ops."""
input_shape = array_ops.shape(inputs)
crop_size = array_ops.stack(
[input_shape[0], self.height, self.width, input_shape[3]])
check = control_flow_ops.Assert(
math_ops.reduce_all(input_shape >= crop_size),
[self.height, self.width])
input_shape = control_flow_ops.with_dependencies([check], input_shape)
limit = input_shape - crop_size + 1
offset = stateless_random_ops.stateless_random_uniform(
array_ops.shape(input_shape),
dtype=crop_size.dtype,
maxval=crop_size.dtype.max,
seed=self._rng.make_seeds()[:, 0]) % limit
return array_ops.slice(inputs, offset, crop_size)
def testDistributionOfStatelessRandomUniform(self):
"""Use Pearson's Chi-squared test to test for uniformity."""
with self.session() as sess, self.test_scope():
for dtype in self._random_types(include_int=True):
seed_t = array_ops.placeholder(dtypes.int32, shape=[2])
n = 1000
maxval = 1
if dtype.is_integer:
maxval = 100
x = stateless.stateless_random_uniform(
shape=[n], seed=seed_t, maxval=maxval, dtype=dtype)
y = sess.run(x, {seed_t: [565656, 121212]})
# Convert y to float and normalize its value to range [0, 1) when
# maxval != 1.
y = y.astype(float) / maxval
# Tests that the values are distributed amongst 10 bins with equal
# probability. 16.92 is the Chi^2 value for 9 degrees of freedom with
# p=0.05. This test is probabilistic and would be flaky if the random
# seed were not fixed.
self.assertLess(random_test_util.chi_squared(y, 10), 16.92)
def testDistributionOfStatelessRandomUniform(self):
"""Use Pearson's Chi-squared test to test for uniformity."""
with self.cached_session() as sess, self.test_scope():
for dtype in self._random_types(include_int=True):
seed_t = array_ops.placeholder(dtypes.int32, shape=[2])
n = 1000
maxval = 1
if dtype.is_integer:
maxval = 100
x = stateless.stateless_random_uniform(
shape=[n], seed=seed_t, maxval=maxval, dtype=dtype)
y = sess.run(x, {seed_t: [565656, 121212]})
if maxval > 1:
# Normalize y to range [0, 1).
y = y.astype(float) / maxval
# Tests that the values are distributed amongst 10 bins with equal
# probability. 16.92 is the Chi^2 value for 9 degrees of freedom with
# p=0.05. This test is probabilistic and would be flaky if the random
# seed were not fixed.
self.assertLess(random_test_util.chi_squared(y, 10), 16.92)
def main(unused_argv):
# Build data and .
print('Loading data.')
x_train, y_train, x_test, y_test = datasets.get_dataset('mnist',
permute_train=True)
# Build the network
init_fn, f, _ = stax.serial(stax.Dense(512, 1., 0.05), stax.Erf(),
stax.Dense(10, 1., 0.05))
key = random.stateless_random_uniform(shape=[2],
seed=[0, 0],
minval=None,
maxval=None,
dtype=np.int32)
_, params = init_fn(key, (1, 784))
# Linearize the network about its initial parameters.
f_lin = nt.linearize(f, params)
# Create and initialize an optimizer for both f and f_lin.
opt_init, opt_apply, get_params = optimizers.momentum(
FLAGS.learning_rate, 0.9)
opt_apply = jit(opt_apply)
state = opt_init(params)
state_lin = opt_init(params)
# Create a cross-entropy loss function.
loss = lambda fx, y_hat: -np.mean(logsoftmax(fx) * y_hat)
# Specialize the loss function to compute gradients for both linearized and
# full networks.
grad_loss = jit(grad(lambda params, x, y: loss(f(params, x), y)))
grad_loss_lin = jit(grad(lambda params, x, y: loss(f_lin(params, x), y)))
# Train the network.
print('Training.')
print('Epoch\tLoss\tLinearized Loss')
print('------------------------------------------')
epoch = 0
steps_per_epoch = 50000 // FLAGS.batch_size
for i, (x, y) in enumerate(
datasets.minibatch(x_train, y_train, FLAGS.batch_size,
FLAGS.train_epochs)):
params = get_params(state)
state = opt_apply(i, grad_loss(params, x, y), state)
params_lin = get_params(state_lin)
state_lin = opt_apply(i, grad_loss_lin(params_lin, x, y), state_lin)
if i % steps_per_epoch == 0:
print('{}\t{}\t{}'.format(epoch, loss(f(params, x), y),
loss(f_lin(params_lin, x), y)))
epoch += 1
# Print out summary data comparing the linear / nonlinear model.
x, y = x_train[:10000], y_train[:10000]
util.print_summary('train', y, f(params, x), f_lin(params_lin, x), loss)
util.print_summary('test', y_test, f(params, x_test),
f_lin(params_lin, x_test), loss)
