def layer_jacobian(self, points, representatives):
"""Computes the Jacobian of the FULLY CONNECTED layer parameters.
Returns (n_points, n_outputs, [weight_shape])
Basically, we assume WLOG that the layer is the first layer then we get
something like for each point:
B_nB_{n-1}...B_1Ax
For each point we can collapse B_n...B_1 into a matrix B of shape
(n_outputs, n_A_outputs)
then we compute the einsum with x to get
(n_outputs, n_A_outputs, n_A_inputs)
which is what we want.
"""
pre_network = Network(self.network.layers[:self.layer_index])
points = pre_network.compute(points)
n_points, A_in_dims = points.shape
representatives = pre_network.compute(representatives)
representatives = self.network.layers[self.layer_index].compute(representatives)
n_points, A_out_dims = representatives.shape
post_network = Network(self.network.layers[(self.layer_index+1):])
# (n_points, A_out, A_out)
jacobian = np.repeat([np.eye(A_out_dims)], n_points, axis=0)
# (A_out, n_points, A_out)
jacobian = jacobian.transpose((1, 0, 2))
for layer in post_network.layers:
if isinstance(layer, LINEAR_LAYERS):
if isinstance(layer, ConcatLayer):
assert not any(isinstance(input_layer, ConcatLayer)
for input_layer in layer.input_layers)
assert all(isinstance(input_layer, LINEAR_LAYERS)
for input_layer in layer.input_layers)
representatives = layer.compute(representatives)
jacobian = jacobian.reshape((A_out_dims * n_points, -1))
jacobian = layer.compute(jacobian, jacobian=True)
jacobian = jacobian.reshape((A_out_dims, n_points, -1))
elif isinstance(layer, ReluLayer):
# (n_points, running_dims)
zero_indices = (representatives <= 0)
representatives[zero_indices] = 0.
# (A_out, n_points, running_dims)
jacobian[:, zero_indices] = 0.
elif isinstance(layer, HardTanhLayer):
big_indices = (representatives >= 1.)
small_indices = (representatives <= -1.)
np.clip(representatives, -1.0, 1.0, out=representatives)
jacobian[:, big_indices] = 0.
jacobian[:, small_indices] = 0.
else:
raise NotImplementedError
# (A_out, n_points, n_classes) -> (n_points, n_classes, A_out)
B = jacobian.transpose((1, 2, 0))
# (n_points, n_classes, n_A_in, n_A_out)
C = np.einsum("nco,ni->ncio", B, points)
return C, B
def linearize_around(self, inputs, network, layer_index):
"""Linearizes a network before/after a certain layer.
We return a 3-tuple, (pre_lins, mid_inputs, post_lins) satisfying:
1) For all x in the same linear region as x_i:
Network(x) = PostLin_i(Layer_l(PreLin_i(x)))
2) For all x_i:
mid_inputs_i = PreLin_i(x_i)
pre_lins and post_lins are lists of FullyConnectedLayers, one per
input, and mid_inputs is a Numpy array.
"""
pre_network = Network(network.layers[:layer_index])
pre_linearized = self.linearize_network(inputs, pre_network)
mid_inputs = pre_network.compute(inputs)
pre_post_inputs = network.layers[layer_index].compute(mid_inputs)
post_network = Network(network.layers[(layer_index + 1):])
post_linearized = self.linearize_network(pre_post_inputs, post_network)
return pre_linearized, mid_inputs, post_linearized
