def randn_sampler_switchover(shape, num_iters, use_gpu=False): # Benchmark by constructing samplers on the threshold of using the randn # rejection sampling and check that this threshold is set correctly by # benchmarking with bounds just above and below this threshold. # The uniform and randn samplers should have about the same performance # at this point. stddev_inside_bounds_before_using_randn = ( _get_stddev_inside_bounds_before_using_randn(use_gpu)) epsilon = 0.001 np.random.seed(1618) # Make it reproducible. # No CSE/CF. optimizer_options = config_pb2.OptimizerOptions( opt_level=config_pb2.OptimizerOptions.L0) config = config_pb2.ConfigProto( graph_options=config_pb2.GraphOptions( optimizer_options=optimizer_options)) with session.Session(config=config) as sess: with ops.device("/cpu:0" if not use_gpu else "/gpu:0"): uniform_sampler_op = control_flow_ops.group( random_ops.parameterized_truncated_normal( shape, means=0., stddevs=1.0, minvals=-stddev_inside_bounds_before_using_randn + epsilon, maxvals=0.01)) randn_sampler_op = control_flow_ops.group( random_ops.parameterized_truncated_normal( shape, means=0., stddevs=1.0, minvals=-stddev_inside_bounds_before_using_randn - epsilon, maxvals=0.01)) # Burn-in to avoid session setup costs in the timing. sess.run(uniform_sampler_op) sess.run(uniform_sampler_op) uniform_dt = timeit.timeit( lambda: sess.run(uniform_sampler_op), number=num_iters) sess.run(randn_sampler_op) sess.run(randn_sampler_op) randn_dt = timeit.timeit( lambda: sess.run(randn_sampler_op), number=num_iters) return randn_dt, uniform_dt
def testParameterizedTruncatedNormalBatched(self): # TODO(b/112289993): Make this test work with dtype np.float64. for dtype in self._random_types() & {np.float32}: with self.session(): with self.test_scope(): count = 10000000 a = -100. b = 100. mu0 = 0. mu1 = 1. sigma = .1 x = random_ops.parameterized_truncated_normal( shape=[2, count], dtype=dtype, means=[mu0, mu1], stddevs=sigma, minvals=[a], maxvals=[b]) self._checkTruncatedNormalIsInRange(x[0], a=a, b=b, mu=mu0, sigma=sigma, count=count, stat_test=True) self._checkTruncatedNormalIsInRange(x[1], a=a, b=b, mu=mu1, sigma=sigma, count=count, stat_test=True)
def validateKolmogorovSmirnov(self, shape, mean, stddev, minval, maxval, seed=1618): try: import scipy.stats # pylint: disable=g-import-not-at-top random_seed.set_random_seed(seed) with self.test_session(use_gpu=True): samples = random_ops.parameterized_truncated_normal(shape, mean, stddev, minval, maxval).eval() assert (~np.isnan(samples)).all() minval = max(mean - stddev * 10, minval) maxval = min(mean + stddev * 10, maxval) dist = scipy.stats.norm(loc=mean, scale=stddev) cdf_min = dist.cdf(minval) cdf_max = dist.cdf(maxval) def truncated_cdf(x): return np.clip((dist.cdf(x) - cdf_min) / (cdf_max - cdf_min), 0.0, 1.0) pvalue = scipy.stats.kstest(samples, truncated_cdf)[1] self.assertGreater(pvalue, 1e-10) except ImportError as e: tf_logging.warn("Cannot test truncated normal op: %s" % str(e))
def testSamplingWithSmallStdDevFarFromBound(self): sample_op = random_ops.parameterized_truncated_normal( shape=(int(1e5), ), means=0.8, stddevs=0.05, minvals=-1., maxvals=1.) new_seed = random_ops.random_uniform([2], seed=1234, minval=0, maxval=(2**31 - 1), dtype=np.int32) sample_op_stateless = stateless.stateless_parameterized_truncated_normal( shape=(int(1e5), ), seed=new_seed, means=0.8, stddevs=0.05, minvals=-1., maxvals=1.) with self.session() as sess: samples, samples_stateless = sess.run( [sample_op, sample_op_stateless]) # 0. is more than 16 standard deviations from the mean, and # should have a likelihood < 1e-57. assert (~np.isnan(samples)).all() assert (~np.isnan(samples_stateless)).all() self.assertAllGreater(samples, 0.) self.assertAllGreater(samples_stateless, 0.)
def testShapeTypes(self): for shape_dtype in [np.int32, np.int64]: shape = np.array([1000], dtype=shape_dtype) sample_op = random_ops.parameterized_truncated_normal(shape=shape, means=0.0, stddevs=0.1, minvals=-1., maxvals=1.) new_seed = random_ops.random_uniform([2], seed=1234, minval=0, maxval=(2**31 - 1), dtype=np.int32) sample_op_stateless = stateless.stateless_parameterized_truncated_normal( shape=shape, seed=new_seed, means=0.0, stddevs=0.1, minvals=-1., maxvals=1.) samples = self.evaluate(sample_op) stateless_samples = self.evaluate(sample_op_stateless) self.assertAllEqual(samples.shape, shape) self.assertAllEqual(stateless_samples.shape, shape)
def validateKolmogorovSmirnov(self, shape, mean, stddev, minval, maxval, seed=1618): try: import scipy.stats # pylint: disable=g-import-not-at-top tf.set_random_seed(seed) with self.test_session(use_gpu=self.use_gpu): samples = random_ops.parameterized_truncated_normal( shape, mean, stddev, minval, maxval).eval() minval = max(mean - stddev * 10, minval) maxval = min(mean + stddev * 10, maxval) dist = scipy.stats.norm(loc=mean, scale=stddev) cdf_min = dist.cdf(minval) cdf_max = dist.cdf(maxval) def truncated_cdf(x): return np.clip((dist.cdf(x) - cdf_min) / (cdf_max - cdf_min), 0.0, 1.0) pvalue = scipy.stats.kstest(samples, truncated_cdf)[1] self.assertGreater(pvalue, 1e-10) except ImportError as e: tf.logging.warn("Cannot test truncated normal op: %s" % str(e))
def _sample_n(self, n, seed=None): sample_and_batch_shape = tf.concat([[n], self.batch_shape_tensor()], 0) flat_batch_and_sample_shape = tf.stack([ tf.reduce_prod(self.batch_shape_tensor()), n]) # In order to be reparameterizable we sample on the truncated_normal of # unit variance and mean and scale (but with the standardized # truncation bounds). std_samples = random_ops.parameterized_truncated_normal( shape=flat_batch_and_sample_shape, means=0.0, stddevs=1.0, minvals=tf.reshape(self._standardized_low, [-1]), maxvals=tf.reshape(self._standardized_high, [-1]), dtype=self.dtype, seed=seed) # The returned shape is [flat_batch x n] std_samples = tf.transpose(std_samples, [1, 0]) std_samples = tf.reshape(std_samples, sample_and_batch_shape) samples = (std_samples * tf.expand_dims(self._scale, axis=0) + tf.expand_dims(self._loc, axis=0)) return samples
def _sample_n(self, n, seed=None): sample_and_batch_shape = tf.concat([[n], self.batch_shape_tensor()], 0) flat_batch_and_sample_shape = tf.stack( [tf.reduce_prod(self.batch_shape_tensor()), n]) # In order to be reparameterizable we sample on the truncated_normal of # unit variance and mean and scale (but with the standardized # truncation bounds). std_samples = random_ops.parameterized_truncated_normal( shape=flat_batch_and_sample_shape, means=0.0, stddevs=1.0, minvals=tf.reshape(self._standardized_low, [-1]), maxvals=tf.reshape(self._standardized_high, [-1]), dtype=self.dtype, seed=seed) # The returned shape is [flat_batch x n] std_samples = tf.transpose(std_samples, [1, 0]) std_samples = tf.reshape(std_samples, sample_and_batch_shape) samples = (std_samples * tf.expand_dims(self._scale, axis=0) + tf.expand_dims(self._loc, axis=0)) return samples
def _std_samples_with_gradients(lower, upper): """Standard truncated Normal with gradient support for low, high.""" # Note: Unlike the convention in tf_probability, # parameterized_truncated_normal returns a tensor with the final dimension # being the sample dimension. std_samples = random_ops.parameterized_truncated_normal( shape=flat_batch_and_sample_shape, means=0.0, stddevs=1.0, minvals=lower, maxvals=upper, dtype=self.dtype, seed=seed) def grad(dy): """Computes a derivative for the min and max parameters. This function implements the derivative wrt the truncation bounds, which get blocked by the sampler. We use a custom expression for numerical stability instead of automatic differentiation on CDF for implicit gradients. Args: dy: output gradients Returns: The standard normal samples and the gradients wrt the upper bound and lower bound. """ # std_samples has an extra dimension (the sample dimension), expand # lower and upper so they broadcast along this dimension. # See note above regarding parameterized_truncated_normal, the sample # dimension is the final dimension. lower_broadcast = lower[..., tf.newaxis] upper_broadcast = upper[..., tf.newaxis] cdf_samples = ((special_math.ndtr(std_samples) - special_math.ndtr(lower_broadcast)) / (special_math.ndtr(upper_broadcast) - special_math.ndtr(lower_broadcast))) # tiny, eps are tolerance parameters to ensure we stay away from giving # a zero arg to the log CDF expression. tiny = np.finfo(self.dtype.as_numpy_dtype).tiny eps = np.finfo(self.dtype.as_numpy_dtype).eps cdf_samples = tf.clip_by_value(cdf_samples, tiny, 1 - eps) du = tf.exp(0.5 * (std_samples**2 - upper_broadcast**2) + tf.log(cdf_samples)) dl = tf.exp(0.5 * (std_samples**2 - lower_broadcast**2) + tf.log1p(-cdf_samples)) # Reduce the gradient across the samples grad_u = tf.reduce_sum(dy * du, axis=-1) grad_l = tf.reduce_sum(dy * dl, axis=-1) return [grad_l, grad_u] return std_samples, grad
def _sample_n(self, n, seed=None): shape = array_ops.concat([[n], self.batch_shape_tensor()], 0) samples = parameterized_truncated_normal(shape, means=self.loc, stddevs=self.scale, minvals=self.minval, maxvals=self.maxval, dtype=self.loc.dtype, seed=seed) return samples
def testSamplingWithSmallStdDevFarFromBound(self): sample_op = random_ops.parameterized_truncated_normal( shape=(int(1e5),), means=0.8, stddevs=0.05, minvals=-1., maxvals=1.) with self.session(use_gpu=True) as sess: samples = sess.run(sample_op) # 0. is more than 16 standard deviations from the mean, and # should have a likelihood < 1e-57. assert (~np.isnan(samples)).all() no_neg_samples = np.sum(samples < 0.) self.assertEqual(no_neg_samples, 0.)
def parameterized_vs_naive(shape, num_iters): np.random.seed(1618) # Make it reproducible. # No CSE/CF. optimizer_options = tf.OptimizerOptions(opt_level=tf.OptimizerOptions.L0) config = tf.ConfigProto(graph_options=tf.GraphOptions(optimizer_options=optimizer_options)) with tf.Session(config=config) as sess: param_op = tf.group(random_ops.parameterized_truncated_normal(shape)) naive_op = tf.group(random_ops.truncated_normal(shape)) param_dt = timeit.timeit(lambda: sess.run(param_op), number=num_iters) naive_dt = timeit.timeit(lambda: sess.run(naive_op), number=num_iters) return param_dt, naive_dt
def parameterized_vs_naive(shape, num_iters): np.random.seed(1618) # Make it reproducible. # No CSE/CF. optimizer_options = tf.OptimizerOptions(opt_level=tf.OptimizerOptions.L0) config = tf.ConfigProto( graph_options=tf.GraphOptions(optimizer_options=optimizer_options)) with tf.Session(config=config) as sess: param_op = tf.group(random_ops.parameterized_truncated_normal(shape)) naive_op = tf.group(random_ops.truncated_normal(shape)) param_dt = timeit.timeit(lambda: sess.run(param_op), number=num_iters) naive_dt = timeit.timeit(lambda: sess.run(naive_op), number=num_iters) return param_dt, naive_dt
def _implParameterizedTruncatedNormalIsInRange(self, a, b, mu, sigma, count, stat_test): # TODO(b/34339814): make this test work with 16 bit float types. for dtype in self._random_types() & {np.float32, np.float64}: with self.session(): with self.test_scope(): x = random_ops.parameterized_truncated_normal( shape=[count], dtype=dtype, means=mu, stddevs=sigma, minvals=a, maxvals=b) self._checkTruncatedNormalIsInRange( x, a=a, b=b, mu=mu, sigma=sigma, count=count, stat_test=stat_test)
def validateMoments(self, shape, mean, stddev, minval, maxval, seed=1618): try: # TruncatedNormalMoments requires scipy.stats. # Give up early if we are unable to import it. import scipy.stats # pylint: disable=g-import-not-at-top,unused-variable tf.set_random_seed(seed) with self.test_session(use_gpu=self.use_gpu): samples = random_ops.parameterized_truncated_normal(shape, mean, stddev, minval, maxval).eval() moments = calculate_moments(samples, self.max_moment) expected_moments = TruncatedNormalMoments(mean, stddev, minval, maxval) num_samples = functools.reduce(lambda x, y: x * y, shape, 1) for i in range(1, len(moments)): self.assertLess(z_test(moments, expected_moments, i, num_samples), self.z_limit) except ImportError as e: tf.logging.warn("Cannot test truncated normal op: %s" % str(e))
def testSamplingWithSmallStdDevFarFromBound(self): sample_op = random_ops.parameterized_truncated_normal( shape=(int(1e5), ), means=0.8, stddevs=0.05, minvals=-1., maxvals=1.) with self.test_session(use_gpu=True) as sess: samples = sess.run(sample_op) # 0. is more than 16 standard deviations from the mean, and # should have a likelihood < 1e-57. # TODO(jjhunt) Sampler is still numerically unstable in this case, # numbers less than 0 should never observed. no_neg_samples = np.sum(samples < 0.) self.assertLess(no_neg_samples, 2.)
def validateMoments(self, shape, mean, stddev, minval, maxval, seed=1618): try: # TruncatedNormalMoments requires scipy.stats. # Give up early if we are unable to import it. import scipy.stats # pylint: disable=g-import-not-at-top,unused-variable tf.set_random_seed(seed) with self.test_session(use_gpu=self.use_gpu): samples = random_ops.parameterized_truncated_normal( shape, mean, stddev, minval, maxval).eval() moments = calculate_moments(samples, self.max_moment) expected_moments = TruncatedNormalMoments(mean, stddev, minval, maxval) num_samples = functools.reduce(lambda x, y: x * y, shape, 1) for i in range(1, len(moments)): self.assertLess( z_test(moments, expected_moments, i, num_samples), self.z_limit) except ImportError as e: tf.logging.warn("Cannot test truncated normal op: %s" % str(e))
def _std_samples_with_gradients(lower, upper): """Standard truncated Normal with gradient support for low, high.""" std_samples = random_ops.parameterized_truncated_normal( shape=flat_batch_and_sample_shape, means=0.0, stddevs=1.0, minvals=lower, maxvals=upper, dtype=self.dtype, seed=seed) def grad(dy): """Computes a derivative for the min and max parameters. This function implements the derivative wrt the truncation bounds, which get blocked by the sampler. We use a custom expression for numerical stability instead of automatic differentiation on CDF for implicit gradients. Args: dy: output gradients Returns: The standard normal samples and the gradients wrt the upper bound and lower bound. """ cdf_samples = ((special_math.ndtr(std_samples) - special_math.ndtr(lower)) / (special_math.ndtr(upper) - special_math.ndtr(lower))) # tiny, eps are tolerance parameters to ensure we stay away from giving # a zero arg to the log CDF expression. tiny = np.finfo(self.dtype.as_numpy_dtype).tiny eps = np.finfo(self.dtype.as_numpy_dtype).eps cdf_samples = tf.clip_by_value(cdf_samples, tiny, 1 - eps) du = tf.exp(0.5 * (std_samples**2 - upper**2) + tf.log(cdf_samples)) dl = tf.exp(0.5 * (std_samples**2 - lower**2) + tf.log(1 - cdf_samples)) # Reduce the gradient across the samples grad_u = tf.reduce_sum(dy * du, axis=-1) grad_l = tf.reduce_sum(dy * dl, axis=-1) return [grad_l, grad_u] return std_samples, grad
def testParameterizedTruncatedNormalBroadcasting(self): for dtype in self._random_types() & {np.float32, np.float64}: with self.session(): with self.test_scope(): a = -1. b = 1. mu = 0. sigma = 1. count = 10000000 x = random_ops.parameterized_truncated_normal( shape=[1, count], dtype=dtype, means=mu, stddevs=sigma, minvals=[a], maxvals=[b]) self._checkTruncatedNormalIsInRange( x, a=a, b=b, mu=mu, sigma=sigma, count=count, stat_test=True)
def parameterized_vs_naive(shape, num_iters, use_gpu=False): np.random.seed(1618) # Make it reproducible. # No CSE/CF. optimizer_options = config_pb2.OptimizerOptions(opt_level=config_pb2.OptimizerOptions.L0) config = config_pb2.ConfigProto(graph_options=config_pb2.GraphOptions(optimizer_options=optimizer_options)) with session.Session(config=config) as sess: with ops.device("/cpu:0" if not use_gpu else None): param_op = control_flow_ops.group(random_ops.parameterized_truncated_normal(shape)) naive_op = control_flow_ops.group(random_ops.truncated_normal(shape)) # Burn-in to avoid session setup costs in the timing. sess.run(param_op) sess.run(param_op) param_dt = timeit.timeit(lambda: sess.run(param_op), number=num_iters) sess.run(naive_op) sess.run(naive_op) naive_dt = timeit.timeit(lambda: sess.run(naive_op), number=num_iters) return param_dt, naive_dt
def validateKolmogorovSmirnov(self, shape, mean, stddev, minval, maxval, use_stateless=False, seed=1618): try: import scipy.stats # pylint: disable=g-import-not-at-top random_seed.set_random_seed(seed) with self.cached_session(): if use_stateless: new_seed = random_ops.random_uniform([2], seed=seed, minval=0, maxval=(2**31 - 1), dtype=np.int32) samples = stateless.stateless_parameterized_truncated_normal( shape, new_seed, mean, stddev, minval, maxval).eval() else: samples = random_ops.parameterized_truncated_normal( shape, mean, stddev, minval, maxval).eval() assert (~np.isnan(samples)).all() minval = max(mean - stddev * 10, minval) maxval = min(mean + stddev * 10, maxval) dist = scipy.stats.norm(loc=mean, scale=stddev) cdf_min = dist.cdf(minval) cdf_max = dist.cdf(maxval) def truncated_cdf(x): return np.clip((dist.cdf(x) - cdf_min) / (cdf_max - cdf_min), 0.0, 1.0) pvalue = scipy.stats.kstest(samples, truncated_cdf)[1] self.assertGreater(pvalue, 1e-10) except ImportError as e: tf_logging.warn("Cannot test truncated normal op: %s" % str(e))
def parameterized_vs_naive(shape, num_iters, use_gpu=False): np.random.seed(1618) # Make it reproducible. # No CSE/CF. optimizer_options = tf.OptimizerOptions(opt_level=tf.OptimizerOptions.L0) config = tf.ConfigProto(graph_options=tf.GraphOptions( optimizer_options=optimizer_options)) with tf.Session(config=config) as sess: with tf.device("/cpu:0" if not use_gpu else None): param_op = tf.group( random_ops.parameterized_truncated_normal(shape)) naive_op = tf.group(random_ops.truncated_normal(shape)) # Burn-in to avoid session setup costs in the timing. sess.run(param_op) sess.run(param_op) param_dt = timeit.timeit(lambda: sess.run(param_op), number=num_iters) sess.run(naive_op) sess.run(naive_op) naive_dt = timeit.timeit(lambda: sess.run(naive_op), number=num_iters) return param_dt, naive_dt
def validateMoments(self, shape, mean, stddev, minval, maxval, use_stateless=False, seed=1618): try: # TruncatedNormalMoments requires scipy.stats. # Give up early if we are unable to import it. random_seed.set_random_seed(seed) with self.cached_session(): if use_stateless: # Generate a seed that stateless ops can use. new_seed = random_ops.random_uniform([2], seed=seed, minval=0, maxval=(2**31 - 1), dtype=np.int32) samples = stateless.stateless_parameterized_truncated_normal( shape, new_seed, mean, stddev, minval, maxval).eval() else: samples = random_ops.parameterized_truncated_normal( shape, mean, stddev, minval, maxval).eval() assert (~np.isnan(samples)).all() moments = calculate_moments(samples, self.max_moment) expected_moments = TruncatedNormalMoments(mean, stddev, minval, maxval) num_samples = functools.reduce(lambda x, y: x * y, shape, 1) for i in range(1, len(moments)): self.assertLess( z_test(moments, expected_moments, i, num_samples), self.z_limit) except ImportError as e: tf_logging.warn("Cannot test truncated normal op: %s" % str(e))