388 lines
15 KiB
Python
388 lines
15 KiB
Python
## @package gradient_checker
|
|
# Module caffe2.python.gradient_checker
|
|
|
|
|
|
|
|
|
|
|
|
import os
|
|
import numpy as np
|
|
|
|
from caffe2.python import core, workspace, net_drawer
|
|
from caffe2.proto import caffe2_pb2
|
|
|
|
|
|
def getGradientForOp(op):
|
|
return core.GradientRegistry.GetGradientForOp(
|
|
op, [s + '_grad' for s in op.output])
|
|
|
|
|
|
def _get_grad_blob(grad_map, input_to_check):
|
|
grad_blob = grad_map[input_to_check]
|
|
|
|
if isinstance(grad_blob, core.BlobReference):
|
|
return workspace.blobs[grad_blob]
|
|
|
|
# If grad_blob is not a single blob, it should be a gradient slice.
|
|
# To make it comparable with the estimiated gradient which is dense,
|
|
# we need to first convert grad_blob to dense gradient.
|
|
assert isinstance(grad_blob, core.GradientSlice)
|
|
dense_grad = 'tmp_dense_grad'
|
|
sparse_to_dense_op = core.CreateOperator(
|
|
'SparseToDense',
|
|
[grad_blob.indices, grad_blob.values, input_to_check],
|
|
dense_grad,
|
|
)
|
|
workspace.RunOperatorOnce(sparse_to_dense_op)
|
|
return workspace.blobs[dense_grad]
|
|
|
|
|
|
def _get_grad(net, outputs, outputs_with_grad, input_values, inputs_with_grads):
|
|
grad_net = net.Clone(net.Name() + "_copy")
|
|
grad_map = grad_net.AddGradientOperators(outputs_with_grad)
|
|
|
|
for name, value in (input_values or {}).items():
|
|
workspace.blobs[name] = value
|
|
|
|
for input_to_check in inputs_with_grads:
|
|
assert input_to_check in grad_map, (
|
|
'{} has no gradient, cannot check net gradient.'.format(
|
|
input_to_check))
|
|
assert str(input_to_check) in workspace.blobs
|
|
|
|
workspace.RunNetOnce(grad_net)
|
|
forward_results = [(output, workspace.blobs[output]) for output in outputs]
|
|
grads = {input_to_check: _get_grad_blob(grad_map, input_to_check)
|
|
for input_to_check in inputs_with_grads}
|
|
|
|
return forward_results, grads, grad_net
|
|
|
|
|
|
def _assert_close(value1, value2, threshold, err_msg=''):
|
|
np.testing.assert_allclose(
|
|
value1, value2,
|
|
atol=threshold, rtol=threshold,
|
|
err_msg=err_msg,
|
|
)
|
|
|
|
delta = np.abs(value1 - value2).flatten()
|
|
return np.mean(delta), max(delta)
|
|
|
|
|
|
class NetGradientChecker(object):
|
|
@staticmethod
|
|
def CompareNets(nets, outputs, outputs_with_grad_ids,
|
|
inputs_with_grads, input_values=None,
|
|
threshold=0.0000001, print_net_images=False):
|
|
def _get_output_with_grad_names(net_outputs):
|
|
return [net_outputs[i] for i in outputs_with_grad_ids]
|
|
|
|
if print_net_images:
|
|
for i, net in enumerate(nets):
|
|
png = net_drawer.GetPydotGraph(net).create_png()
|
|
with open("caffe2_net_forward_" + str(i) + net.Name() + ".png",
|
|
'wb') \
|
|
as f:
|
|
f.write(png)
|
|
|
|
results = [
|
|
_get_grad(net, net_outputs,
|
|
_get_output_with_grad_names(net_outputs),
|
|
input_values, inputs_with_grads)
|
|
for net, net_outputs in zip(nets, outputs)
|
|
]
|
|
|
|
if print_net_images:
|
|
_, _, backward_nets = zip(*results)
|
|
for i, net in enumerate(backward_nets):
|
|
png = net_drawer.GetPydotGraph(net).create_png()
|
|
with open("caffe2_net_" + str(i) + net.Name() + ".png", 'wb') \
|
|
as f:
|
|
f.write(png)
|
|
|
|
first_net_results, first_net_grads, _ = results[0]
|
|
for net_results, net_grads, _ in results[1:]:
|
|
assert len(net_results) == len(first_net_results)
|
|
for idx, ((blob1, blob_value1), (blob2, blob_value2)) in enumerate(
|
|
zip(first_net_results, net_results)):
|
|
_assert_close(
|
|
blob_value1, blob_value2, threshold,
|
|
err_msg="Different forward pass results for output id {}. "
|
|
"Corresponding output blobs: {} and {}".format(
|
|
idx, blob1, blob2))
|
|
|
|
assert net_grads.keys() == first_net_grads.keys()
|
|
for blob, blob_grad_value in net_grads.items():
|
|
_assert_close(
|
|
first_net_grads[blob], blob_grad_value, threshold,
|
|
err_msg="Different gradients for input {}".format(blob))
|
|
|
|
@staticmethod
|
|
def Check(net, outputs_with_grad, input_values,
|
|
input_to_check, step_size=0.0001,
|
|
threshold=0.05, print_net=True):
|
|
|
|
net_results, net_grads, full_net = _get_grad(
|
|
net, [], outputs_with_grad, input_values, [input_to_check])
|
|
analytic_grad = net_grads[input_to_check]
|
|
|
|
def GetLoss(new_value):
|
|
workspace.blobs[input_to_check] = new_value
|
|
workspace.RunNetOnce(full_net)
|
|
return sum([
|
|
workspace.blobs[output]
|
|
for output in outputs_with_grad
|
|
]).sum()
|
|
|
|
def GetValue(dim, delta):
|
|
input_value = input_values[input_to_check].copy()
|
|
input_value.flat[dim] += delta
|
|
return input_value
|
|
|
|
grad_estimate = np.zeros_like(input_values[input_to_check])
|
|
for dim in range(input_values[input_to_check].size):
|
|
pos_loss = GetLoss(GetValue(dim, step_size))
|
|
neg_loss = GetLoss(GetValue(dim, -step_size))
|
|
grad_estimate.flat[dim] = (pos_loss - neg_loss) / step_size / 2
|
|
|
|
err_msg = "Error in gradient check for net_copy {}".format(
|
|
net.Name())
|
|
if print_net:
|
|
err_msg += ": {}".format(net.Proto())
|
|
|
|
return _assert_close(analytic_grad, grad_estimate, threshold, err_msg)
|
|
|
|
|
|
class GradientChecker:
|
|
"""A gradient checker in Python.
|
|
|
|
This is not the most efficient way to check gradients, as the Python
|
|
interface will involve a lot of copies back and forth operations. Use at your
|
|
own risk.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
stepsize,
|
|
threshold,
|
|
device_option=None,
|
|
workspace_name="gradient_check",
|
|
input_device_options=None,
|
|
):
|
|
self._stepsize = stepsize
|
|
self._threshold = threshold
|
|
self._device_option = device_option or caffe2_pb2.DeviceOption()
|
|
self._workspace_name = workspace_name
|
|
if input_device_options is None:
|
|
self._input_device_options = {}
|
|
else:
|
|
self._input_device_options = input_device_options
|
|
|
|
def GetLossAndGrad(
|
|
self, op, grad_ops, inputs, input_names, input_to_check, grad_name,
|
|
outputs_with_grads
|
|
):
|
|
for i in range(len(inputs)):
|
|
workspace.FeedBlob(input_names[i], inputs[i],
|
|
self._input_device_options.get(
|
|
input_names[i], self._device_option))
|
|
x = inputs[input_to_check]
|
|
# Run.
|
|
workspace.RunOperatorOnce(op)
|
|
loss = 0.
|
|
# Get Loss and feed in the gradients, run gradient ops.
|
|
for idx in outputs_with_grads:
|
|
name = op.output[idx]
|
|
arr = workspace.FetchBlob(name)
|
|
loss += (arr**2).sum()
|
|
workspace.FeedBlob(name + '_grad', arr, self._device_option)
|
|
loss /= 2.
|
|
# Run gradient ops
|
|
workspace.RunOperatorsOnce(grad_ops)
|
|
# Get gradients
|
|
if isinstance(grad_name, core.GradientSlice):
|
|
workspace.FeedBlob('zeros', np.zeros_like(x, dtype=np.float32))
|
|
workspace.FeedBlob('ones', np.ones(1, dtype=np.float32))
|
|
gv_cpu_op = core.CreateOperator(
|
|
'EnsureCPUOutput', grad_name.values, grad_name.values + '_cpu',
|
|
device_option=self._device_option
|
|
)
|
|
gi_cpu_op = core.CreateOperator(
|
|
'EnsureCPUOutput', grad_name.indices, grad_name.indices + '_cpu',
|
|
device_option=self._device_option
|
|
)
|
|
sparse_to_dense_op = core.CreateOperator(
|
|
'ScatterWeightedSum',
|
|
[
|
|
'zeros', 'ones', grad_name.indices + '_cpu',
|
|
grad_name.values + '_cpu', 'ones'
|
|
],
|
|
'zeros',
|
|
)
|
|
workspace.RunOperatorOnce(gv_cpu_op)
|
|
workspace.RunOperatorOnce(gi_cpu_op)
|
|
workspace.RunOperatorOnce(sparse_to_dense_op)
|
|
grad = workspace.FetchBlob('zeros')
|
|
else:
|
|
grad = workspace.FetchBlob(grad_name)
|
|
return loss, grad
|
|
|
|
def CheckSimple(
|
|
self,
|
|
op,
|
|
inputs,
|
|
input_to_check,
|
|
outputs_with_grads,
|
|
grad_ops=None,
|
|
input_device_options=None,
|
|
ensure_outputs_are_inferred=False,
|
|
):
|
|
"""Checks the operator in a very simple fashion by stacking a sum of
|
|
squares on the top.
|
|
|
|
Inputs:
|
|
op: the operator to be checked.
|
|
inputs: the input data in numpy arrays.
|
|
input_to_check: an index specifying which input blob we should
|
|
check.
|
|
outputs_with_grads: indices specifying which output blobs will we
|
|
need to check gradients with. For these outputs, we will collect a
|
|
squared sum and also feed in their gradients.
|
|
grad_operator: the gradient operator. If not given, we will get the
|
|
gradient operator from the gradient registry.
|
|
input_device_options: an optional mapping from input names to
|
|
DeviceOptions (to override the default DeviceOption)
|
|
ensure_outputs_are_inferred: if set will assert that the gradient output
|
|
shapes matches the inferred shapes
|
|
Outputs:
|
|
boolean: True if it passes, False if it does not pass.
|
|
"""
|
|
# Entering the checker workspace
|
|
old_ws_name = workspace.CurrentWorkspace()
|
|
if self._workspace_name != old_ws_name:
|
|
workspace.SwitchWorkspace(self._workspace_name, True)
|
|
|
|
op.device_option.CopyFrom(self._device_option)
|
|
if grad_ops is None:
|
|
# TODO(jiayq): use the gradient registration instead of the old
|
|
# hack.
|
|
grad_ops, g_input = getGradientForOp(op)
|
|
|
|
|
|
_input_device_options = input_device_options or \
|
|
core.InferOpBlobDevicesAsDict(op)[0]
|
|
# First, feed in the input.
|
|
for i, arr in enumerate(inputs):
|
|
workspace.FeedBlob(
|
|
op.input[i], arr,
|
|
_input_device_options.get(
|
|
op.input[i], self._device_option))
|
|
|
|
# Get the loss and gradient for the original.
|
|
grad_name = g_input[input_to_check]
|
|
loss, grad = self.GetLossAndGrad(
|
|
op, grad_ops, inputs, op.input, input_to_check, grad_name,
|
|
outputs_with_grads,
|
|
)
|
|
grad_estimate = np.zeros_like(inputs[input_to_check])
|
|
if grad_estimate.shape != grad.shape:
|
|
raise Exception(
|
|
"Mismatched gradient shapes: estimated ({}), grad ({})".format(
|
|
grad_estimate.shape, grad.shape))
|
|
|
|
if ensure_outputs_are_inferred:
|
|
self._assertInferTensorChecks(op, grad_ops)
|
|
|
|
full_grad_check = os.getenv('CAFFE2_FULL_GRAD_CHECK') == '1'
|
|
|
|
dims_to_check = inputs[input_to_check].size
|
|
for current_dim in range(dims_to_check):
|
|
# Grad check is very expensive (as it involves running the op from
|
|
# scratch for each of the input tensor elements). Thus, let's
|
|
# run it by default only on a small subset of dimensions. Here we
|
|
# apply very scientific approach: the first and the last 3 elements
|
|
# of each tensor. Pass CAFFE2_FULL_GRAD_CHECK=1 env var to enable
|
|
# the full check
|
|
if not full_grad_check and current_dim >= 3 and \
|
|
current_dim + 3 < dims_to_check:
|
|
grad_estimate.flat[current_dim] = grad.flat[current_dim]
|
|
continue
|
|
# Positive gradient
|
|
inputs[input_to_check].flat[current_dim] += self._stepsize
|
|
pos_loss, _ = self.GetLossAndGrad(
|
|
op, grad_ops, inputs, op.input, input_to_check, grad_name,
|
|
outputs_with_grads
|
|
)
|
|
# Negative gradient
|
|
inputs[input_to_check].flat[current_dim] -= self._stepsize * 2
|
|
neg_loss, _ = self.GetLossAndGrad(
|
|
op, grad_ops, inputs, op.input, input_to_check, grad_name,
|
|
outputs_with_grads
|
|
)
|
|
# Recover the value
|
|
inputs[input_to_check].flat[current_dim] += self._stepsize
|
|
grad_estimate.flat[current_dim] = (
|
|
pos_loss - neg_loss) / self._stepsize / 2
|
|
# Now, check correctness
|
|
fail_mat = ~np.isclose(
|
|
grad, grad_estimate, atol=self._threshold, rtol=self._threshold)
|
|
if np.any(fail_mat):
|
|
idx = np.flatnonzero(fail_mat)
|
|
print('Failed. [idx, grad, grad_estimate] are:')
|
|
print(np.vstack([idx, grad.flat[idx], grad_estimate.flat[idx]]).T)
|
|
ret = False
|
|
else:
|
|
ret = True
|
|
# After finishing, cleaning up things.
|
|
if self._workspace_name != old_ws_name:
|
|
# We reset the workspace to make sure everything intermediate is
|
|
# cleaned up. Note that there is no need to delete a workspace -
|
|
# when empty it takes a very limited amount of memory.
|
|
workspace.ResetWorkspace()
|
|
workspace.SwitchWorkspace(old_ws_name)
|
|
return ret, grad, grad_estimate
|
|
|
|
def _assertInferTensorChecks(self, op, grad_ops):
|
|
tmp_net = caffe2_pb2.NetDef()
|
|
tmp_net.op.extend([op])
|
|
tmp_net.op.extend(grad_ops)
|
|
inferred_shapes, inferred_types = workspace.InferShapesAndTypes(
|
|
[tmp_net],
|
|
nets_proto=True,
|
|
)
|
|
|
|
outputs = set()
|
|
for grad_op in grad_ops:
|
|
outputs.update(grad_op.output)
|
|
|
|
for output in outputs:
|
|
if output not in inferred_shapes:
|
|
raise Exception(
|
|
"expected output {} to be inferred".format(output))
|
|
blob = workspace.FetchBlob(output)
|
|
correct_shape = list(blob.shape)
|
|
inferred_shape = list(inferred_shapes[output])
|
|
if correct_shape != inferred_shape:
|
|
raise Exception(
|
|
"Mismatched inferred shape: want({}), got({})".format(
|
|
correct_shape, inferred_shape))
|
|
|
|
if type(blob) is np.ndarray:
|
|
if blob.dtype == np.dtype('float64'):
|
|
correct_type = caffe2_pb2.TensorProto.DOUBLE
|
|
elif blob.dtype == np.dtype('float32'):
|
|
correct_type = caffe2_pb2.TensorProto.FLOAT
|
|
elif blob.dtype == np.dtype('int32'):
|
|
correct_type = caffe2_pb2.TensorProto.INT32
|
|
elif blob.dtype == np.dtype('int64'):
|
|
correct_type = caffe2_pb2.TensorProto.INT64
|
|
else:
|
|
correct_type = "unknown {}".format(np.dtype)
|
|
else:
|
|
correct_type = str(type(blob))
|
|
inferred_type = inferred_types[output]
|
|
if correct_type != inferred_type:
|
|
raise Exception(
|
|
"Mismatched inferred type: want({}), got({})".format(
|
|
correct_type, inferred_type))
|