def _interpolate_in_direction(self, vec, steps):
"""Interpolate the loss in the given direction.
Args:
vec: Direction to move weights in.
steps: Array of coefficients to multiply vector by.
Returns:
Array of [steps, losses], where the loss is evaluated at each step along
the vec direction.
"""
# Make all vector manipulations on CPU so we don't run out of memory
with tf.device("/cpu:0"):
orig_weights = self.model.get_weights()
flat_orig_weights = tfutils.flatten_tensor_list(orig_weights)
losses = []
# TODO Again assuming that computing last-layer vectors
expanded_vec = np.zeros(flat_orig_weights.shape)
expanded_vec[-vec.shape[0]:] = vec
for alpha in steps:
target_weights = flat_orig_weights + alpha * expanded_vec
unflatten_target_weights = K.get_session().run(
tfutils.unflatten_tensor_list(
target_weights,
self.model.trainable_weights))
self.model.set_weights(unflatten_target_weights)
loss = tfutils.compute_sample_mean_tensor(
self.model, self.train_batches, self.model.total_loss)
losses.append(loss)
self.model.set_weights(orig_weights)
return np.array([steps, losses])
def _compute_metrics(self, batches, logs, prefix):
means = tfutils.compute_sample_mean_tensor(self.model, batches,
self.all_tensors)
self.record_scalar(prefix + 'loss', means[0])
self.record_scalar(prefix + 'acc', means[1])
logs[prefix + 'loss'] = means[0]
logs[prefix + 'acc'] = means[1]
def test_compute_sample_mean_tensor(self):
K.clear_session()
d = 12
n = 20
batch_size = n // 4
x = np.random.rand(n, d).astype(np.float32)
y = np.sin(2 * np.pi * x[:, 0]).reshape((-1, 1)).astype(np.float32)
# Linear regression
model = keras.models.Sequential()
model.add(keras.layers.Dense(1, use_bias=False, input_shape=(d,)))
model.add(keras.layers.Activation('relu'))
model.add(keras.layers.Dense(10))
model.add(keras.layers.Activation('relu'))
model.add(keras.layers.Dense(1))
model.compile(loss='mean_squared_error', optimizer=keras.optimizers.SGD())
tfutils.keras_compute_tensors(model, x, y, model.total_loss)
grad_t = tfutils.flatten_tensor_list(
tf.gradients(model.total_loss, model.trainable_weights))
grad = tfutils.keras_compute_tensors(model, x, y, grad_t)
batches = tfutils.MiniBatchMaker(x, y, batch_size)
actual_grad = tfutils.compute_sample_mean_tensor(model, batches, grad_t)
self.assertTrue(np.allclose(grad, actual_grad))
def _compute_Hv_overlap(self, batches, v):
Hv = tfutils.compute_sample_mean_tensor(self.model, batches, self.Hv,
{self.v: v})
v_norm = np.linalg.norm(v)
Hv_norm = np.linalg.norm(Hv)
Hv_dot_v = np.dot(Hv, v)
Hv_v_overlap = Hv_dot_v / v_norm / Hv_norm
Hv_eigenvalue = Hv_dot_v / v_norm**2
return Hv, Hv_v_overlap, Hv_eigenvalue