def _extrapolation_is_linear(self, float_dtype):
"""Check that extrapolation is linear with respect to the endpoint knots.
Generate random splines and query them outside of the support of the
spline, and veify that extrapolation is linear with respect to the
endpoint knots.
Args:
float_dtype: the dtype of the floats to be tested.
"""
n = 256
# Generate queries above and below the support of the spline.
x_below = float_dtype(-(np.random.uniform(size=1024)) * (n - 1))
x_above = float_dtype((np.random.uniform(size=1024) + 1.) * (n - 1))
for _ in range(8):
values = float_dtype(np.random.normal(size=n))
tangents = float_dtype(np.random.normal(size=n))
# Query the spline below its support and check that it's a linear ramp
# with the slope and bias of the beginning of the spline.
y_below = cubic_spline.interpolate1d(x_below, values, tangents)
y_below_true = tangents[0] * x_below + values[0]
self.assertAllClose(y_below, y_below_true)
# Query the spline above its support and check that it's a linear ramp
# with the slope and bias of the end of the spline.
y_above = cubic_spline.interpolate1d(x_above, values, tangents)
y_above_true = tangents[-1] * (x_above - (n - 1)) + values[-1]
self.assertAllClose(y_above, y_above_true)
def _interpolate1d(self, x, values, tangents):
"""Compute interpolate1d(x, values, tangents) and its derivative.
This is just a helper function around cubic_spline.interpolate1d() that does
the necessary work to get derivatives and handle TensorFlow sessions.
Args:
x: A np.array of values to interpolate with.
values: A np.array of knot values for the spline.
tangents: A np.array of knot tangents for the spline.
Returns:
A tuple containing:
(An np.array of interpolated values,
A np.array of derivatives of interpolated values wrt `x`)
Typical usage example:
y, dy_dx = self._interpolate1d(x, values, tangents)
"""
x = tf.convert_to_tensor(x)
with tf.GradientTape() as tape:
tape.watch(x)
y = cubic_spline.interpolate1d(x, values, tangents)
dy_dx = tape.gradient(y, x)
return y, dy_dx
def log_base_partition_function(self, alpha):
r"""Approximate the distribution's log-partition function with a 1D spline.
Because the partition function (Z(\alpha) in the paper) of the distribution
is difficult to model analytically, we approximate it with a (transformed)
cubic hermite spline: Each alpha is pushed through a nonlinearity before
being used to interpolate into a spline, which allows us to use a relatively
small spline to accurately model the log partition function over the range
of all non-negative input values.
Args:
alpha: A tensor or scalar of single or double precision floats containing
the set of alphas for which we would like an approximate log partition
function. Must be non-negative, as the partition function is undefined
when alpha < 0.
Returns:
An approximation of log(Z(alpha)) accurate to within 1e-6
"""
float_dtype = alpha.dtype
# The partition function is undefined when `alpha`< 0.
assert_ops = [tf.Assert(tf.reduce_all(alpha >= 0.), [alpha])]
with tf.control_dependencies(assert_ops):
# Transform `alpha` to the form expected by the spline.
x = partition_spline_curve(alpha)
# Interpolate into the spline.
return cubic_spline.interpolate1d(
x * tf.cast(self._spline_x_scale, float_dtype),
tf.cast(self._spline_values, float_dtype),
tf.cast(self._spline_tangents, float_dtype))
def _interpolate1d(self, x, values, tangents):
"""Compute interpolate1d(x, values, tangents) and its derivative.
This is just a helper function around cubic_spline.interpolate1d() that does
the necessary work to get derivatives and handle TensorFlow sessions.
Args:
x: A np.array of values to interpolate with.
values: A np.array of knot values for the spline.
tangents: A np.array of knot tangents for the spline.
Returns:
A tuple containing:
(An np.array of interpolated values,
A np.array of derivatives of interpolated values wrt `x`)
Typical usage example:
y, dy_dx = self._interpolate1d(x, values, tangents)
"""
with tf.Session() as sess:
x_ph = tf.placeholder(x.dtype, x.shape)
values_ph = tf.placeholder(values.dtype, values.shape)
tangents_ph = tf.placeholder(tangents.dtype, tangents.shape)
y_ph = cubic_spline.interpolate1d(x_ph, values_ph, tangents_ph)
y, (dy_dx, ) = sess.run(
(y_ph, tf.gradients(tf.reduce_sum(y_ph), (x_ph))), {
x_ph: x,
values_ph: values,
tangents_ph: tangents,
})
return y, dy_dx
def _interpolation_reproduces_values_at_knots(self, float_dtype): """Check that interpolating at a knot produces the value at that knot.""" n = 32768 x = np.arange(n, dtype=float_dtype) values = float_dtype(np.random.normal(size=n)) tangents = float_dtype(np.random.normal(size=n)) y = cubic_spline.interpolate1d(x, values, tangents) self.assertAllClose(y, values)
def _interpolation_preserves_dtype(self, float_dtype): """Check that interpolating at a knot produces the value at that knot.""" n = 16 x = float_dtype(np.random.normal(size=n)) values = float_dtype(np.random.normal(size=n)) tangents = float_dtype(np.random.normal(size=n)) y = cubic_spline.interpolate1d(x, values, tangents) self.assertDTypeEqual(y, float_dtype)
def _linear_ramps_reproduce_correctly(self, float_dtype):
"""Check that interpolating a ramp reproduces a ramp.
Generate linear ramps, render them into splines, and then interpolate and
extrapolate the splines and verify that they reproduce the ramp.
Args:
float_dtype: the dtype of the floats to be tested.
"""
n = 256
# Generate queries inside and outside the support of the spline.
x = float_dtype((np.random.uniform(size=1024) * 2 - 0.5) * (n - 1))
idx = np.arange(n, dtype=float_dtype)
for _ in range(8):
slope = np.random.normal()
bias = np.random.normal()
values = slope * idx + bias
tangents = np.ones_like(values) * slope
y = cubic_spline.interpolate1d(x, values, tangents)
y_true = slope * x + bias
self.assertAllClose(y, y_true)
def log_base_partition_function(alpha):
r"""Approximate the distribution's log-partition function with a 1D spline.
Because the partition function (Z(\alpha) in the paper) of the distribution is
difficult to model analytically, we approximate it with a (transformed) cubic
hermite spline: Each alpha is pushed through a nonlinearity before being used
to interpolate into a spline, which allows us to use a relatively small spline
to accurately model the log partition function over the range of all
non-negative input values.
Args:
alpha: A tensor or scalar of single or double precision floats containing
the set of alphas for which we would like an approximate log partition
function. Must be non-negative, as the partition function is undefined
when alpha < 0.
Returns:
An approximation of log(Z(alpha)) accurate to within 1e-6
"""
float_dtype = alpha.dtype
# Load the values, tangents, and x-coordinate scaling of a spline that
# approximates the partition function. This was produced by running
# the script in fit_partition_spline.py
with util.get_resource_as_file(
'robust_loss/data/partition_spline.npz') as spline_file:
with np.load(spline_file, allow_pickle=False) as f:
x_scale = tf.cast(f['x_scale'], float_dtype)
values = tf.cast(f['values'], float_dtype)
tangents = tf.cast(f['tangents'], float_dtype)
# The partition function is undefined when `alpha`< 0.
assert_ops = [tf.Assert(tf.reduce_all(alpha >= 0.), [alpha])]
with tf.control_dependencies(assert_ops):
# Transform `alpha` to the form expected by the spline.
x = partition_spline_curve(alpha)
# Interpolate into the spline.
return cubic_spline.interpolate1d(x * x_scale, values, tangents)
def main(argv):
if len(argv) > 1:
raise app.UsageError('Too many command-line arguments.')
# Parameters governing how the x coordinate of the spline will be laid out.
# We will construct a spline with knots at
# [0 : 1 / x_scale : x_max],
# by fitting it to values sampled at
# [0 : 1 / (x_scale * redundancy) : x_max]
x_max = 12
x_scale = 1024
redundancy = 4 # Must be >= 2 for the spline to be useful.
spline_spacing = 1. / (x_scale * redundancy)
x_knots = np.arange(
0, x_max + spline_spacing, spline_spacing, dtype=np.float64)
table = []
# We iterate over knots, and for each knot recover the alpha value
# corresponding to that knot with inv_partition_spline_curve(), and then
# with that alpha we accurately approximate its partition function using
# numerical_base_partition_function().
for x_knot in x_knots:
alpha = distribution.inv_partition_spline_curve(x_knot).numpy()
partition = numerical_base_partition_function(alpha).numpy()
table.append((x_knot, alpha, partition))
print(table[-1])
table = np.array(table)
x = table[:, 0]
alpha = table[:, 1]
y_gt = np.log(table[:, 2])
# We grab the values from the true log-partition table that correpond to
# knots, by looking for where x * x_scale is an integer.
mask = np.abs(np.round(x * x_scale) - (x * x_scale)) <= 1e-8
values = y_gt[mask]
# Initialize `tangents` using a central differencing scheme.
values_pad = np.concatenate([[values[0] - values[1] + values[0]], values,
[values[-1] - values[-2] + values[-1]]], 0)
tangents = (values_pad[2:] - values_pad[:-2]) / 2.
# Construct the spline's value and tangent TF variables, constraining the last
# knot to have a fixed value Z(infinity) and a tangent of zero.
n = len(values)
tangents = tf.Variable(tangents, tf.float64)
values = tf.Variable(values, tf.float64)
# Fit the spline.
num_iters = 10001
optimizer = tf.keras.optimizers.SGD(learning_rate=1e-9, momentum=0.99)
trace = []
for ii in range(num_iters):
with tf.GradientTape() as tape:
tape.watch([values, tangents])
# Fix the endpoint to be a known constant with a zero tangent.
i_values = tf.where(
np.arange(n) == (n - 1),
tf.ones_like(values) * 0.70526025442689566, values)
i_tangents = tf.where(
np.arange(n) == (n - 1), tf.zeros_like(tangents), tangents)
i_y = cubic_spline.interpolate1d(x * x_scale, i_values, i_tangents)
# We minimize the maximum residual, which makes for a very ugly
# optimization problem but works well in practice.
i_loss = tf.reduce_max(tf.abs(i_y - y_gt))
grads = tape.gradient(i_loss, [values, tangents])
optimizer.apply_gradients(zip(grads, [values, tangents]))
trace.append(i_loss.numpy())
if (ii % 200) == 0:
print('{:5d}: {:e}'.format(ii, trace[-1]))
mask = alpha <= 4
max_error_a4 = np.max(np.abs(i_y[mask] - y_gt[mask]))
max_error = np.max(np.abs(i_y - y_gt))
print('Max Error (a <= 4): {:e}'.format(max_error_a4))
print('Max Error: {:e}'.format(max_error))
# Just a sanity-check on the error.
assert max_error_a4 <= 5e-7
assert max_error <= 5e-7
# Save the spline to disk.
np.savez(
'./data/partition_spline.npz',
x_scale=x_scale,
values=i_values.numpy(),
tangents=i_tangents.numpy())
def main(argv):
if len(argv) > 1:
raise app.UsageError('Too many command-line arguments.')
# Parameters governing how the x coordinate of the spline will be laid out.
# We will construct a spline with knots at
# [0 : 1 / x_scale : x_max],
# by fitting it to values sampled at
# [0 : 1 / (x_scale * redundancy) : x_max]
x_max = 12
x_scale = 1024
redundancy = 4 # Must be >= 2 for the spline to be useful.
spline_spacing = 1. / (x_scale * redundancy)
x_knots = np.arange(0,
x_max + spline_spacing,
spline_spacing,
dtype=np.float64)
table = []
with tf.Session() as sess:
x_knot_ph = tf.placeholder(dtype=tf.float64, shape=())
alpha_ph = distribution.inv_partition_spline_curve(x_knot_ph)
partition_ph = numerical_base_partition_function(alpha_ph)
# We iterate over knots, and for each knot recover the alpha value
# corresponding to that knot with inv_partition_spline_curve(), and then
# with that alpha we accurately approximate its partition function using
# numerical_base_partition_function().
for x_knot in x_knots:
alpha, partition = sess.run((alpha_ph, partition_ph),
{x_knot_ph: x_knot})
table.append((x_knot, alpha, partition))
print(table[-1])
table = np.array(table)
x = table[:, 0]
alpha = table[:, 1]
y_gt = np.log(table[:, 2])
# We grab the values from the true log-partition table that correpond to
# knots, by looking for where x * x_scale is an integer.
mask = np.abs(np.round(x * x_scale) - (x * x_scale)) <= 1e-8
values = y_gt[mask]
# Initialize `tangents` using a central differencing scheme.
values_pad = np.concatenate([[values[0] - values[1] + values[0]], values,
[values[-1] - values[-2] + values[-1]]], 0)
tangents = (values_pad[2:] - values_pad[:-2]) / 2.
# Construct the spline's value and tangent TF variables, constraining the last
# knot to have a fixed value Z(infinity) and a tangent of zero.
n = len(values)
tangents = tf.Variable(tangents, tf.float64)
tangents = tf.where(
np.arange(n) == (n - 1), tf.zeros_like(tangents), tangents)
values = tf.Variable(values, tf.float64)
values = tf.where(
np.arange(n) == (n - 1),
tf.ones_like(tangents) * 0.70526025442689566, values)
# Interpolate into the spline.
y = cubic_spline.interpolate1d(x * x_scale, values, tangents)
# We minimize the maximum residual, which makes for a very ugly optimization
# problem but appears to work in practice, and is what we most care about.
loss = tf.reduce_max(tf.abs(y - y_gt))
# Fit the spline.
num_iters = 10001
with tf.Session() as sess:
global_step = tf.Variable(0, trainable=False)
opt = tf.train.MomentumOptimizer(learning_rate=1e-9, momentum=0.99)
step = opt.minimize(loss, global_step=global_step)
sess.run(tf.global_variables_initializer())
trace = []
for ii in range(num_iters):
_, i_loss, i_values, i_tangents, i_y = sess.run(
[step, loss, values, tangents, y])
trace.append(i_loss)
if (ii % 200) == 0:
print('%5d: %e' % (ii, i_loss))
mask = alpha <= 4
print('Max Error (a <= 4): %e' % np.max(np.abs(i_y[mask] - y_gt[mask])))
print('Max Error: %e' % np.max(np.abs(i_y - y_gt)))
# Save the spline to disk.
np.savez('./data/partition_spline.npz',
x_scale=x_scale,
values=i_values,
tangents=i_tangents)
