# Copyright 2015 The TensorFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== """Gradient checker for any ops, graphs. The gradient checker verifies numerically that an op/graph properly computes the gradients """ from __future__ import absolute_import from __future__ import division from __future__ import print_function import numpy as np from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops from tensorflow.python.ops import gradients from tensorflow.python.ops import math_ops from tensorflow.python.platform import tf_logging as logging from tensorflow.python.util.tf_export import tf_export def _product(t): if isinstance(t, int): return t else: y = 1 for x in t: y *= x return y def _extra_feeds(extra_feed_dict, new_feeds): if not extra_feed_dict: return new_feeds r = {} r.update(extra_feed_dict) r.update(new_feeds) return r def _compute_theoretical_jacobian(x, x_shape, x_data, dy, dy_shape, dx, extra_feed_dict): """Computes the theoretical Jacobian for dy/dx. Computes the theoretical Jacobian using the ops generated by compute_gradient(). Args: x: the tensor "x". x_shape: the dimensions of x as a tuple or an array of ints. x_data: a numpy parray as the input data for x dy: the tensor "dy". dy_shape: the dimensions of dy as a tuple or an array of ints. dx: Tensor or IndexedSlices representing dx extra_feed_dict: dict that allows fixing specified tensor values during the jacobian calculation. Returns: A 2-d numpy array representing the Jacobian for dy/dx. It has "x_size" rows and "dy_size" columns where "x_size" is the number of elements in x and "dy_size" is the number of elements in dy. Raises: ValueError: If `dy` is empty but the gradient is nonzero. """ # Complex vectors are treated as vectors of twice as many reals. if x.dtype.is_complex: x_shape = tuple(x_shape) + (2,) dy_factor = 2 if dy.dtype.is_complex else 1 # To compute the jacobian, we treat x and y as one-dimensional vectors. x_size = _product(x_shape) x_val_size = _product(x_shape[1:]) # This is used for sparse gradients dy_size = _product(dy_shape) * dy_factor # Allocate 2-D Jacobian, with x dimensions smashed into the first # dimension and y dimensions smashed into the second. jacobian = np.zeros((x_size, dy_size), dtype=x.dtype.real_dtype.as_numpy_dtype) # For each of the entry of dy, we set this to be 1 and # everything else to be 0 and compute the backprop -- this will give us one # one column of the Jacobian matrix. dy_data = np.zeros(dy_shape, dtype=dy.dtype.as_numpy_dtype) dy_data_flat = dy_data.ravel().view(dy.dtype.real_dtype.as_numpy_dtype) sess = ops.get_default_session() for col in range(dy_size): dy_data_flat[col] = 1 if isinstance(dx, ops.IndexedSlices): backprop_indices, backprop_values = sess.run( [dx.indices, dx.values], feed_dict=_extra_feeds(extra_feed_dict, {x: x_data, dy: dy_data})) for i, v in zip(backprop_indices, backprop_values): r_begin = i * x_val_size r_end = r_begin + x_val_size jacobian[r_begin:r_end, col] += v.flat else: assert isinstance(dx, ops.Tensor), "dx = " + str(dx) backprop = sess.run( dx, feed_dict=_extra_feeds(extra_feed_dict, {x: x_data, dy: dy_data})) jacobian[:, col] = backprop.ravel().view(jacobian.dtype) dy_data_flat[col] = 0 # If the output is empty, run the gradients at least once and make sure # they produce zeros. if not dy_size: backprop = sess.run( dx, feed_dict=_extra_feeds(extra_feed_dict, {x: x_data, dy: dy_data})) if backprop.shape != x_data.shape: raise ValueError("Empty gradient has wrong shape: expected %s, got %s" % (x_data.shape, backprop.shape)) if np.any(backprop): raise ValueError("Empty tensor with nonzero gradients") logging.vlog(1, "Theoretical Jacobian =\n%s", jacobian) return jacobian def _compute_numeric_jacobian(x, x_shape, x_data, y, y_shape, delta, extra_feed_dict): """Computes the numeric Jacobian for dy/dx. Computes the numeric Jacobian by slightly perturbing the inputs and measuring the differences on the output. Args: x: the tensor "x". x_shape: the dimensions of x as a tuple or an array of ints. x_data: a numpy array as the input data for x y: the tensor "y". y_shape: the dimensions of y as a tuple or an array of ints. delta: the amount of perturbation we give to the input extra_feed_dict: dict that allows fixing specified tensor values during the jacobian calculation. Returns: A 2-d numpy array representing the Jacobian for dy/dx. It has "x_size" rows and "y_size" columns where "x_size" is the number of elements in x and "y_size" is the number of elements in y. """ # bfloat16 doesn't have enough bits to represent high precision numbers such # as delta. Convert to float32 here. Since numeric_jacobian is expected to # be the groundtruth to compare against, it shouldn't lose any information. if x.dtype == dtypes.bfloat16: x = math_ops.cast(x, dtypes.float32) if y.dtype == dtypes.bfloat16: y = math_ops.cast(y, dtypes.float32) if x_data.dtype == dtypes.bfloat16.as_numpy_dtype: x_data = x_data.astype(np.float32) # To compute the jacobian, we treat x and y as one-dimensional vectors x_size = _product(x_shape) * (2 if x.dtype.is_complex else 1) y_size = _product(y_shape) * (2 if y.dtype.is_complex else 1) x_dtype = x.dtype.real_dtype.as_numpy_dtype y_dtype = y.dtype.real_dtype.as_numpy_dtype # Make sure we have the right types x_data = np.asarray(x_data, dtype=x.dtype.as_numpy_dtype) scale = np.asarray(2 * delta, dtype=y_dtype)[()] jacobian = np.zeros((x_size, y_size), dtype=x_dtype) # For each of the entry of x, we slightly perturbs this by adding and # subtracting a delta and then compute difference between the outputs. This # will give us one row of the Jacobian matrix. for row in range(x_size): x_pos = x_data.copy() x_neg = x_data.copy() x_pos.ravel().view(x_dtype)[row] += delta y_pos = y.eval(feed_dict=_extra_feeds(extra_feed_dict, {x: x_pos})) x_neg.ravel().view(x_dtype)[row] -= delta y_neg = y.eval(feed_dict=_extra_feeds(extra_feed_dict, {x: x_neg})) diff = (y_pos - y_neg) / scale jacobian[row, :] = diff.ravel().view(y_dtype) logging.vlog(1, "Numeric Jacobian =\n%s", jacobian) return jacobian def _compute_dx_and_dy(x, y, y_shape): """Returns a node to compute gradient of y wrt x.""" # We make up a dy so that we can compute the gradients. We don't really use # the value of dy -- we will always feed it. We need to add an identity node # so that we can always feed it properly. Otherwise, for the Add operation, # dx is the same as dy and we cannot fetch the tensor that we are feeding. with x.graph.as_default(): dy_orig = constant_op.constant(1.0, shape=y_shape, dtype=y.dtype) dy = array_ops.identity(dy_orig) # We compute the gradients for y wrt. x grads = gradients.gradients(y, x, dy) assert len(grads) == 1 return grads[0], dy_orig def _compute_gradient(x, x_shape, dx, y, y_shape, dy, x_init_value=None, delta=1e-3, extra_feed_dict=None): """Computes the theoretical and numerical jacobian.""" t = dtypes.as_dtype(x.dtype) allowed_types = [dtypes.float16, dtypes.bfloat16, dtypes.float32, dtypes.float64, dtypes.complex64, dtypes.complex128] assert t.base_dtype in allowed_types, "Don't support type %s for x" % t.name t2 = dtypes.as_dtype(y.dtype) assert t2.base_dtype in allowed_types, "Don't support type %s for y" % t2.name if x_init_value is not None: i_shape = list(x_init_value.shape) assert(list(x_shape) == i_shape), "x_shape = %s, init_data shape = %s" % ( x_shape, i_shape) x_data = x_init_value else: x_data = np.random.random_sample(x_shape).astype(t.as_numpy_dtype) if t.is_complex: x_data.imag = np.random.random_sample(x_shape) jacob_t = _compute_theoretical_jacobian( x, x_shape, x_data, dy, y_shape, dx, extra_feed_dict=extra_feed_dict) jacob_n = _compute_numeric_jacobian( x, x_shape, x_data, y, y_shape, delta, extra_feed_dict=extra_feed_dict) return jacob_t, jacob_n def _compute_gradient_list(x, x_shape, y, y_shape, x_init_value=None, delta=1e-3, init_targets=None, extra_feed_dict=None): """Compute gradients for a list of x values.""" assert isinstance(x, list) dx, dy = zip(*[_compute_dx_and_dy(xi, y, y_shape) for xi in x]) if init_targets is not None: assert isinstance(init_targets, (list, tuple)) for init in init_targets: init.run() if x_init_value is None: x_init_value = [None] * len(x) ret = [_compute_gradient(xi, x_shapei, dxi, y, y_shape, dyi, x_init_valuei, delta, extra_feed_dict=extra_feed_dict) for xi, x_shapei, dxi, dyi, x_init_valuei in zip(x, x_shape, dx, dy, x_init_value)] return ret @tf_export("test.compute_gradient") def compute_gradient(x, x_shape, y, y_shape, x_init_value=None, delta=1e-3, init_targets=None, extra_feed_dict=None): """Computes and returns the theoretical and numerical Jacobian. If `x` or `y` is complex, the Jacobian will still be real but the corresponding Jacobian dimension(s) will be twice as large. This is required even if both input and output is complex since TensorFlow graphs are not necessarily holomorphic, and may have gradients not expressible as complex numbers. For example, if `x` is complex with shape `[m]` and `y` is complex with shape `[n]`, each Jacobian `J` will have shape `[m * 2, n * 2]` with J[::2, ::2] = d(Re y)/d(Re x) J[::2, 1::2] = d(Im y)/d(Re x) J[1::2, ::2] = d(Re y)/d(Im x) J[1::2, 1::2] = d(Im y)/d(Im x) Args: x: a tensor or list of tensors x_shape: the dimensions of x as a tuple or an array of ints. If x is a list, then this is the list of shapes. y: a tensor y_shape: the dimensions of y as a tuple or an array of ints. x_init_value: (optional) a numpy array of the same shape as "x" representing the initial value of x. If x is a list, this should be a list of numpy arrays. If this is none, the function will pick a random tensor as the initial value. delta: (optional) the amount of perturbation. init_targets: list of targets to run to initialize model params. TODO(mrry): remove this argument. extra_feed_dict: dict that allows fixing specified tensor values during the Jacobian calculation. Returns: Two 2-d numpy arrays representing the theoretical and numerical Jacobian for dy/dx. Each has "x_size" rows and "y_size" columns where "x_size" is the number of elements in x and "y_size" is the number of elements in y. If x is a list, returns a list of two numpy arrays. """ if extra_feed_dict is None: extra_feed_dict = {} if isinstance(x, list): return _compute_gradient_list(x, x_shape, y, y_shape, x_init_value, delta, init_targets, extra_feed_dict=extra_feed_dict) else: if init_targets is not None: assert isinstance(init_targets, (list, tuple)) for init in init_targets: init.run() dx, dy = _compute_dx_and_dy(x, y, y_shape) ret = _compute_gradient(x, x_shape, dx, y, y_shape, dy, x_init_value, delta, extra_feed_dict=extra_feed_dict) return ret @tf_export("test.compute_gradient_error") def compute_gradient_error(x, x_shape, y, y_shape, x_init_value=None, delta=1e-3, init_targets=None, extra_feed_dict=None): """Computes the gradient error. Computes the maximum error for dy/dx between the computed Jacobian and the numerically estimated Jacobian. This function will modify the tensors passed in as it adds more operations and hence changing the consumers of the operations of the input tensors. This function adds operations to the current session. To compute the error using a particular device, such as a GPU, use the standard methods for setting a device (e.g. using with sess.graph.device() or setting a device function in the session constructor). Args: x: a tensor or list of tensors x_shape: the dimensions of x as a tuple or an array of ints. If x is a list, then this is the list of shapes. y: a tensor y_shape: the dimensions of y as a tuple or an array of ints. x_init_value: (optional) a numpy array of the same shape as "x" representing the initial value of x. If x is a list, this should be a list of numpy arrays. If this is none, the function will pick a random tensor as the initial value. delta: (optional) the amount of perturbation. init_targets: list of targets to run to initialize model params. extra_feed_dict: dict that allows fixing specified tensor values during the Jacobian calculation. Returns: The maximum error in between the two Jacobians. """ grad = compute_gradient(x, x_shape, y, y_shape, x_init_value, delta, init_targets, extra_feed_dict=extra_feed_dict) if isinstance(grad, tuple): grad = [grad] error = 0 for j_t, j_n in grad: if j_t.size or j_n.size: # Handle zero size tensors correctly error = np.maximum(error, np.fabs(j_t - j_n).max()) return error