def reorder_spatial_axes(tensor):
"""
Assumes we are getting a C, H, N, or C, H, W, N, or C, D, H, W, N
"""
spatial_axes = tensor.axes.spatial_axes()
batch_axes = tensor.axes.batch_axes()
if len(spatial_axes) == 0 or len(spatial_axes) > 3:
raise ValueError(
'spatial ops can only operate on tensors with 1, 2, or 3 spatial axes.'
'Found {}'.format(spatial_axes))
if not batch_axes:
raise ValueError('spatial ops require a batch axis')
if not tensor.axes.channel_axis():
c = ng.make_axis(length=1, name='C')
tensor = ng.expand_dims(tensor, c, 0)
channel_axes = ng.make_axes(tensor.axes.channel_axis())
if len(spatial_axes) == 1:
w = ng.make_axis(length=1, name=_WIDTH)
tensor = ng.expand_dims(tensor, w, 0)
spatial_axes = spatial_axes + w
if len(spatial_axes) == 2:
d = ng.make_axis(length=1, name=_DEPTH)
tensor = ng.expand_dims(tensor, d, 0)
spatial_axes = ng.make_axes([d]) + spatial_axes
new_axes = channel_axes + spatial_axes + batch_axes
return ng.axes_with_order(tensor, new_axes)
def _build_module(self, input_layer):
# Dueling Network
# state value tower - V
output_axis = ng.make_axis(self.num_actions, name='q_values')
state_value = neon.Sequential([
neon.Affine(nout=256,
activation=neon.Rectlin(),
weight_init=self.weights_init,
bias_init=self.biases_init),
neon.Affine(nout=1,
weight_init=self.weights_init,
bias_init=self.biases_init)
])(input_layer)
# action advantage tower - A
action_advantage_unnormalized = neon.Sequential([
neon.Affine(nout=256,
activation=neon.Rectlin(),
weight_init=self.weights_init,
bias_init=self.biases_init),
neon.Affine(axes=output_axis,
weight_init=self.weights_init,
bias_init=self.biases_init)
])(input_layer)
action_advantage = action_advantage_unnormalized - ng.mean(
action_advantage_unnormalized)
repeated_state_value = ng.expand_dims(
ng.slice_along_axis(state_value, state_value.axes[0], 0),
output_axis, 0)
# merge to state-action value function Q
self.output = repeated_state_value + action_advantage
def _build_module(self, input_layer):
# This is almost exactly the same as Dueling Network but we predict the future measurements for each action
multistep_measurements_size = self.measurements_size[0] * self.num_predicted_steps_ahead
# actions expectation tower (expectation stream) - E
with name_scope("expectation_stream"):
expectation_stream = neon.Sequential([
neon.Affine(nout=256, activation=neon.Rectlin(),
weight_init=self.weights_init, bias_init=self.biases_init),
neon.Affine(nout=multistep_measurements_size,
weight_init=self.weights_init, bias_init=self.biases_init)
])(input_layer)
# action fine differences tower (action stream) - A
with name_scope("action_stream"):
action_stream_unnormalized = neon.Sequential([
neon.Affine(nout=256, activation=neon.Rectlin(),
weight_init=self.weights_init, bias_init=self.biases_init),
neon.Affine(nout=self.num_actions * multistep_measurements_size,
weight_init=self.weights_init, bias_init=self.biases_init),
neon.Reshape((self.num_actions, multistep_measurements_size))
])(input_layer)
action_stream = action_stream_unnormalized - ng.mean(action_stream_unnormalized)
repeated_expectation_stream = ng.slice_along_axis(expectation_stream, expectation_stream.axes[0], 0)
repeated_expectation_stream = ng.expand_dims(repeated_expectation_stream, output_axis, 0)
# merge to future measurements predictions
self.output = repeated_expectation_stream + action_stream
def test_conv_flatten_deriv(n4_hw12_c3_5x5):
"""
Test deriv of conv followed by flatten
"""
cf = ConvParams(**n4_hw12_c3_5x5)
axes_rsck = ng.make_axes([cf.ax_f[2], cf.ax_f[3], cf.ax_f[0], cf.ax_f[-1]])
axes_rsck_prime = ng.make_axes([ng.make_axis(name=ax.name + 'p', length=ax.length)
for ax in axes_rsck])
axes_nmpqk = ng.make_axes([cf.ax_o[-1], cf.ax_o[1], cf.ax_o[2], cf.ax_o[3], cf.ax_o[0]])
# broadcast input / filter axes
input_var = ng.variable(cf.ax_i).named('input')
input_val = np.ones(input_var.axes.lengths)
filter_rsck_prime = ng.variable(axes_rsck_prime).named('filter')
filter_var = filter_rsck_prime
filter_rsck = ng.cast_axes(filter_rsck_prime, axes_rsck).named('frsck')
filter_trsck = ng.expand_dims(filter_rsck, cf.ax_f[1], 0).named('ftrsck')
filter_ctrsk = ng.axes_with_order(filter_trsck, axes=cf.ax_f).named('ctrsk')
# convolution
output_kmpqn = ng.convolution(cf.conv_params, input_var, filter_ctrsk, axes=cf.ax_o)
output_nmpqk = ng.axes_with_order(output_kmpqn, axes=axes_nmpqk)
# slice away the oD
out_slicing = [slice(None), 0, slice(None), slice(None), slice(None)]
output_npqk = ng.tensor_slice(output_nmpqk, out_slicing)
output = ng.flatten_at(output_npqk, idx=1)
# cost and grad
cost = ng.sum(output, out_axes=())
filter_val = np.ones(filter_var.axes.lengths)
with ExecutorFactory() as factory:
conv_comp = factory.executor(output, filter_var, input_var)
grad_filter_num_comp = factory.numeric_derivative(cost, filter_var, 1.0, input_var)
grad_filter_sym_comp = factory.derivative(cost, filter_var, input_var)
grad_input_num_comp = factory.numeric_derivative(cost, input_var, 1.0, filter_var)
grad_input_sym_comp = factory.derivative(cost, input_var, filter_var)
conv_val = conv_comp(filter_val, input_val)
conv_val_num = np.empty_like(conv_val)
conv_val_num.fill(np.prod(cf.ax_f.lengths[:-1]))
ng.testing.assert_allclose(conv_val, conv_val_num)
grad_filter_num_val = grad_filter_num_comp(filter_val, input_val)
grad_filter_sym_val = grad_filter_sym_comp(filter_val, input_val)
ng.testing.assert_allclose(grad_filter_num_val, grad_filter_sym_val)
grad_input_num_val = grad_input_num_comp(input_val, filter_val)
grad_input_sym_val = grad_input_sym_comp(input_val, filter_val)
ng.testing.assert_allclose(grad_input_num_val, grad_input_sym_val)
def test_conv_flatten_deriv(transformer_factory):
"""
Test deriv of conv followed by flatten
"""
# set shape
C, D, H, W, N = (3, 1, 28, 28, 8)
C, T, R, S, K = (3, 1, 5, 5, 32)
# i, f, o axes
ax_i = ng.make_axes([ax.C, ax.D, ax.H, ax.W, ax.N])
ax_f = ng.make_axes([ax.C, ax.T, ax.R, ax.S, ax.K])
ax_o = ng.make_axes([
ng.make_axis(32, roles=[ar.Channel]),
ng.make_axis(1, roles=[ar.Depth]),
ng.make_axis(24, roles=[ar.Height]),
ng.make_axis(24, roles=[ar.Width]), ax.N
])
ax_i.set_shape((C, D, H, W, N))
ax_f.set_shape((C, T, R, S, K))
params = dict(pad_d=0, pad_h=0, pad_w=0, str_d=1, str_h=1, str_w=1)
axes_rsck = ng.make_axes([ax.R, ax.S, ax.C, ax.K])
axes_rsck_prime = ng.make_axes(
[ng.make_axis(l) for l in axes_rsck.lengths])
# broadcast input / filter axes
image = ng.constant(np.ones(ax_i.lengths), ax_i)
filter = ng.variable(axes_rsck_prime, initial_value=np.ones((R, S, C, K)))
filter_casted = ng.cast_axes(filter, axes_rsck)
filter_casted = ng.expand_dims(filter_casted, ax.T, 0)
filter_casted = ng.axes_with_order(filter_casted, axes=ax_f)
# convolution
output = ng.convolution(params, image, filter_casted, axes=ax_o)
oC, oD, oH, oW, oN = output.axes
output = ng.axes_with_order(output,
axes=ng.make_axes([oN, oD, oH, oW, oC]))
# slice away the oD
out_slicing = [slice(None), 0, slice(None), slice(None), slice(None)]
conv = ng.Slice(output, out_slicing)
flatten = ng.flatten_at(conv, idx=1)
# cost and grad
cost = ng.sum(flatten, reduction_axes=flatten.axes)
grad = ng.deriv(cost, filter)
# compute
conv_grad_comp = executor([conv, grad])
conv_val, grad_val = conv_grad_comp()
assert np.allclose(conv_val, np.zeros_like(conv_val) + 75.)
assert np.allclose(grad_val, np.zeros_like(grad_val) + 4608.)
def expand_with_name(tensor, axis, index=0):
if isinstance(axis, Axis):
if axis in tensor.axes:
return tensor, axis
if (axis.length is not None) and (axis.length > 1):
raise IncompatibleAxesError("Cannot expand tensor to an axis with length > 1: {}"
", length={}".format(axis.name, axis.length))
axis.length = 1
else:
if axis in tensor.axes.names:
return tensor, tensor.axes.find_by_name(axis)[0]
axis = ng.make_axis(name=axis, length=1)
return ng.expand_dims(tensor, axis, index), axis
def make_reduction_op(ng_op_type, onnx_node, ng_input):
# type: (Callable, NodeWrapper, TensorOp) -> Op
"""
Create an ngraph Op node for a reduction operation (min, max, sum, etc.)
:param ng_op_type: an ngraph reduction factory function such as ng.max, etc.
:param onnx_node: wrapped ONNX node
:param ng_input: ngraph Op to be used as input to the reduction node
"""
reduction_ng_axes = get_reduction_axes(onnx_node)
op = ng_op_type(ng_input, reduction_axes=reduction_ng_axes)
if onnx_node.get_attribute_value('keepdims', default=1):
for axis in reduction_ng_axes:
pos = ng_input.axes.index(axis)
new_axis = ng.make_axis(length=1, name=axis.name)
op = ng.expand_dims(op, new_axis, pos)
return op
def test_shuffled_deriv(transformer_factory):
# This gets the axes of a delta in a generate_add_delta in a different order than the
# value being updated
C = ng.make_axis(length=3)
T = ng.make_axis(length=1)
R = ng.make_axis(length=5)
S = ng.make_axis(length=5)
axes = [R, S, C]
v = ng.variable([ng.make_axis(_.length) for _ in axes])
rsc = ng.cast_axes(v, axes)
trsc = ng.expand_dims(rsc, T, 0)
ctrs = ng.axes_with_order(trsc, axes=[C, T, R, S])
cost = ng.sum(ctrs, out_axes=None)
grad = ng.deriv(cost, v)
with ExecutorFactory() as ex:
d_fun = ex.executor(grad)
d_fun()
def test_shuffled_deriv():
# This gets the axes of a delta in a generate_add_delta in a different order than the
# value being updated
ax = ng.make_name_scope("ax")
ax.C = ng.make_axis(3)
ax.T = ng.make_axis(1)
ax.R = ng.make_axis(5)
ax.S = ng.make_axis(5)
axes = [ax.R, ax.S, ax.C]
v = ng.variable([ng.make_axis(_.length) for _ in axes])
rsc = ng.cast_axes(v, axes)
trsc = ng.expand_dims(rsc, ax.T, 0)
ctrs = ng.axes_with_order(trsc, axes=[ax.C, ax.T, ax.R, ax.S])
cost = ng.sum(ctrs, out_axes=None)
grad = ng.deriv(cost, v)
ex = ExecutorFactory()
d_fun = ex.executor(grad)
d_fun()
def MaxPool(self, tf_node, inputs):
"""
Performs the max pooling on the input.
Arguments:
tf_node: NodeDef object, the tensorflow node to convert.
inputs: List of ngraph Ops as inputs to this node.
Returns:
A ngraph Op corresponding to the tensorflow node.
Inputs to tf_node:
input
TODO: assume default tensorflow layout NHWC, RSCK,
need to support NCHW as well
need to clean up / merge with conv2d
Axes:
Tensorflow Ngraph
in (N, H, W, C) (C, D, H, W, N)
out (N, P, Q, K) (K, M, P, Q, N)
Notes on output shape:
https://www.tensorflow.org/api_docs/python/nn.html#convolution
"""
image = inputs[0]
# TODO: currently NHWC only
assert tf_node.attr['data_format'].s.decode("ascii") == "NHWC"
# new axes
C, D, H, W, K, M, P, Q = [ng.make_axis() for _ in range(8)]
N = ng.make_axis(name='N')
D.length, M.length = 1, 1 # only supports 2D conv for now
# tf's input axes
ax_i_tf = ng.make_axes([N, H, W, C])
ax_i_tf.set_shape(image.axes.lengths)
# ksize params
tf_ksize = [int(s) for s in list(tf_node.attr['ksize'].list.i)]
if len(tf_ksize) != 4:
raise ValueError("Length of ksize my be 4.")
if tf_ksize[0] != 1:
raise NotImplementedError('Ksize on batch axis (N) must be 1.')
if tf_ksize[3] != 1:
raise NotImplementedError('Ksize on channel axis (C) must be 1.'
'Cross map pooling to be implemented.')
R_length, S_length = tf_ksize[1:3]
T_length = J_length = 1
# strides params
tf_strides = [int(s) for s in list(tf_node.attr['strides'].list.i)]
if len(tf_strides) != 4:
raise ValueError("Length of strides my be 4.")
if tf_strides[0] != 1:
raise NotImplementedError('Strides on batch axis (N) must be 1.')
if tf_strides[3] != 1:
raise NotImplementedError('Strides on channel axis (C) must be 1.')
str_h, str_w = tf_strides[1], tf_strides[2]
# padding params
padding = tf_node.attr['padding'].s.decode("ascii")
pad_t, pad_b, pad_l, pad_r = common_conv2d_pool_padding(
image.axes.lengths, (R_length, S_length, C.length, C.length),
tf_strides, padding)
if pad_t != pad_b or pad_l != pad_r:
raise NotImplementedError("Requires symmetric padding in ngraph:"
"pad_t(%s) == pad_b(%s) and"
"pad_l(%s) == pad_r(%s)" %
(pad_t, pad_b, pad_l, pad_r))
# pooling params
params = dict(op='max',
pad_d=0, pad_h=pad_t, pad_w=pad_l, pad_c=0,
str_d=1, str_h=str_h, str_w=str_w, str_c=1,
J=J_length, T=T_length, R=R_length, S=S_length)
# tf's output axes
ax_o_tf = ng.make_axes([N, P, Q, K])
ax_o_tf.set_shape(common_conv2d_pool_output_shape(image.axes.lengths,
(R_length, S_length,
C.length, C.length),
tf_strides, padding))
# ngraph's i, f, o axes
ax_i = ng.make_axes([C, D, H, W, N])
ax_o = ng.make_axes([K, M, P, Q, N])
# image NHWC -> CDHWN
image = ng.cast_axes(image, ng.make_axes([N, H, W, C]))
image = ng.expand_dims(image, D, 1) # NHWC -> NDHWC
image = ng.axes_with_order(image, ax_i) # NDHWC -> CDHWN
# pooling
output = ng.pooling(params, image, axes=ax_o)
# output KMPQN -> NPQK
# KMPQN -> NMPQK
output = ng.axes_with_order(output, ng.make_axes(
[N, M, P, Q, K]))
# NMPQK -> NPQK
output = ng.tensor_slice(output, [slice(None), 0, slice(None),
slice(None), slice(None)])
return output
def Conv2D(self, tf_node, inputs):
"""
Computes a 2-D convolution given 4D input and filter tensors.
Arguments:
tf_node: NodeDef object, the tensorflow node to convert.
inputs: List of ngraph Ops as inputs to this node.
Returns:
A ngraph Op corresponding to the tensorflow node.
Inputs to tf_node:
input, filter
TODO: assume default tensorflow layout NHWC, RSCK,
need to support NCHW as well
need to clean up / merge with maxpool
Axes:
Tensorflow Ngraph
in (N, H, W, C) (C, D, H, W, N)
filter (R, S, C, K) (C, T, R, S, K)
out (N, P, Q, K) (K, M, P, Q, N)
Notes on output shape:
https://www.tensorflow.org/api_docs/python/nn.html#convolution
"""
image, weight = inputs
# TODO: currently NHWC only
if tf_node.attr['data_format'].s.decode("ascii") != "NHWC":
raise NotImplementedError("Only supports NHWC import for now.")
# check in_C == f_C
if image.axes.lengths[3] != weight.axes.lengths[2]:
raise ValueError("Image's C dimension (%s) must be equal to "
"filter's C dimension (%s)."
% (image.axes.lengths[3], weight.axes.lengths[2]))
# strides params
tf_strides = [int(s) for s in list(tf_node.attr['strides'].list.i)]
if len(tf_strides) != 4:
raise ValueError("Length of strides my be 4.")
if tf_strides[0] != 1:
raise NotImplementedError('Strides on batch axis (N) must be 1.')
if tf_strides[3] != 1:
raise NotImplementedError('Strides on channel axis (C) must be 1.')
str_h, str_w = tf_strides[1], tf_strides[2]
# padding params
padding = tf_node.attr['padding'].s.decode("ascii")
pad_t, pad_b, pad_l, pad_r = common_conv2d_pool_padding(
image.axes.lengths, weight.axes.lengths, tf_strides, padding)
if pad_t != pad_b or pad_l != pad_r:
raise NotImplementedError("Requires symmetric padding in ngraph:"
"pad_t(%s) == pad_b(%s) and"
"pad_l(%s) == pad_r(%s)" %
(pad_t, pad_b, pad_l, pad_r))
# conv params
params = dict(pad_d=0, pad_h=pad_t, pad_w=pad_l,
str_d=1, str_h=str_h, str_w=str_w,
dil_d=1, dil_h=1, dil_w=1)
# new axes
C, D, H, W, T, R, S, K, M, P, Q = [ng.make_axis() for _ in range(11)]
N = ng.make_axis(name='N')
D.length, T.length, M.length = 1, 1, 1 # only supports 2D conv for now
# tf's i, f, o axes
ax_i_tf = ng.make_axes([N, H, W, C])
ax_f_tf = ng.make_axes([R, S, C, K])
ax_o_tf = ng.make_axes([N, P, Q, K])
ax_i_tf.set_shape(image.axes.lengths)
ax_f_tf.set_shape(weight.axes.lengths)
ax_o_tf.set_shape(common_conv2d_pool_output_shape(image.axes.lengths,
weight.axes.lengths,
tf_strides, padding))
# ngraph's i, f, o axes
ax_i = ng.make_axes([C, D, H, W, N])
ax_f = ng.make_axes([C, T, R, S, K])
ax_o = ng.make_axes([K, M, P, Q, N])
# image NHWC -> CDHWN
image = ng.cast_axes(image, ng.make_axes([N, H, W, C]))
image = ng.expand_dims(image, D, 1) # NHWC -> NDHWC
image = ng.axes_with_order(image, ax_i) # NDHWC -> CDHWN
# weights RSCK -> CTRSK
weight = ng.cast_axes(weight, ng.make_axes([R, S, C, K]))
weight = ng.expand_dims(weight, T, 0) # RSCK -> TRSCK
weight = ng.axes_with_order(weight, ax_f) # TRSCK -> CTRSK
# convolution
output = ng.convolution(params, image, weight, axes=ax_o)
# output KMPQN -> NPQK
# KMPQN -> NMPQK
output = ng.axes_with_order(output, ng.make_axes([N, M, P, Q, K]))
# NMPQK -> NPQK
output = ng.tensor_slice(output, [slice(None), 0, slice(None),
slice(None), slice(None)])
return output
def MaxPool(self, tf_node, inputs):
"""
Performs the max pooling on the input.
Arguments:
tf_node: NodeDef object, the tensorflow node tso convert.
inputs: List of ngraph Ops as inputs to this node.
Returns:
A ngraph Op corresponding to the tensorflow node.
Inputs to tf_node:
input
TODO: assume default tensorflow layout NHWC, RSCK,
need to support NCHW as well
need to clean up / merge with conv2d
Notes on output shape:
https://www.tensorflow.org/api_docs/python/nn.html#convolution
"""
image = inputs[0]
# TODO: currently NHWC only
assert tf_node.attr['data_format'].s.decode("ascii") == "NHWC"
# set axes shape
ax_N = ng.make_axis(batch=True)
ax_C = ng.make_axis(roles=[ar.Channel])
ax_D = ng.make_axis(roles=[ar.Depth])
ax_H = ng.make_axis(roles=[ar.Height])
ax_W = ng.make_axis(roles=[ar.Width])
ng.make_axes([ax_N, ax_H, ax_W, ax_C]).set_shape(image.axes.lengths)
ax_D.length = 1
# ksize params
tf_ksize = [int(s) for s in list(tf_node.attr['ksize'].list.i)]
if len(tf_ksize) != 4:
raise ValueError("Length of ksize my be 4.")
if tf_ksize[0] != 1:
raise NotImplementedError('Ksize on batch axis (N) must be 1.')
if tf_ksize[3] != 1:
raise NotImplementedError('Ksize on channel axis (C) must be 1.'
'Cross map pooling to be implemented.')
R, S = tf_ksize[1:3]
T = J = 1
# strides params
tf_strides = [int(s) for s in list(tf_node.attr['strides'].list.i)]
if len(tf_strides) != 4:
raise ValueError("Length of strides my be 4.")
if tf_strides[0] != 1:
raise NotImplementedError('Strides on batch axis (N) must be 1.')
if tf_strides[3] != 1:
raise NotImplementedError('Strides on channel axis (C) must be 1.')
str_h, str_w = tf_strides[1], tf_strides[2]
# padding params
padding = tf_node.attr['padding'].s.decode("ascii")
pad_t, pad_b, pad_l, pad_r = tf_conv2d_pool_padding(
image.axes.lengths, (R, S, ax_C.length, ax_C.length), tf_strides,
padding)
if pad_t != pad_b or pad_l != pad_r:
raise NotImplementedError("Requires symmetric padding in ngraph:"
"pad_t(%s) == pad_b(%s) and"
"pad_l(%s) == pad_r(%s)" %
(pad_t, pad_b, pad_l, pad_r))
# pooling params
params = dict(op='max',
pad_d=0,
pad_h=pad_t,
pad_w=pad_l,
pad_c=0,
str_d=1,
str_h=str_h,
str_w=str_w,
str_c=1,
J=J,
T=T,
R=R,
S=S)
# i, f, o axes
ax_i = ng.make_axes([ax_C, ax_D, ax_H, ax_W, ax_N])
ax_o = ng.make_axes([
spatial_axis(ax_i, J, params['pad_c'], params['str_c'],
ar.Channel),
spatial_axis(ax_i, T, params['pad_d'], params['str_d'], ar.Depth),
spatial_axis(ax_i, R, params['pad_h'], params['str_h'], ar.Height),
spatial_axis(ax_i, S, params['pad_w'], params['str_w'], ar.Width),
ax_N
])
# broadcast input / filter axes
image = ng.cast_axes(image, ng.make_axes([ax_N, ax_H, ax_W, ax_C]))
image = ng.expand_dims(image, ax_D, 1) # NHWC -> NDHWC
image = ng.axes_with_order(image, axes=ax_i) # NDHWC -> CDHWN
# pooling
output = ng.pooling(params, image, axes=ax_o)
# cast back to NHWC
oC, oD, oH, oW, oN = output.axes
output = ng.broadcast(output, ng.make_axes([oN, oD, oH, oW, oC]))
# slice away the oD
out_slicing = [slice(None), 0, slice(None), slice(None), slice(None)]
output = ng.Slice(output, out_slicing)
return output
def Conv2D(self, tf_node, inputs):
"""
Computes a 2-D convolution given 4D input and filter tensors.
Arguments:
tf_node: NodeDef object, the tensorflow node to convert.
inputs: List of ngraph Ops as inputs to this node.
Returns:
A ngraph Op corresponding to the tensorflow node.
Inputs to tf_node:
input, filter
TODO: assume default tensorflow layout NHWC, RSCK,
need to support NCHW as well
need to clean up / merge with maxpool
Notes on output shape:
https://www.tensorflow.org/api_docs/python/nn.html#convolution
"""
image, weight = inputs
# TODO: currently NHWC only
assert tf_node.attr['data_format'].s.decode("ascii") == "NHWC"
# set axes shape
ax_N = ng.make_axis(batch=True)
ax_C = ng.make_axis(roles=[ar.Channel])
ax_D = ng.make_axis(roles=[ar.Depth])
ax_H = ng.make_axis(roles=[ar.Height])
ax_W = ng.make_axis(roles=[ar.Width])
ax_T = ng.make_axis(roles=[ar.Depth])
ax_R = ng.make_axis(roles=[ar.Height])
ax_S = ng.make_axis(roles=[ar.Width])
ax_K = ng.make_axis(roles=[ar.Channelout])
ng.make_axes([ax_N, ax_H, ax_W, ax_C]).set_shape(image.axes.lengths)
ng.make_axes([ax_R, ax_S, ax_C, ax_K]).set_shape(weight.axes.lengths)
ax_D.length = 1
ax_T.length = 1
# strides params
tf_strides = [int(s) for s in list(tf_node.attr['strides'].list.i)]
if len(tf_strides) != 4:
raise ValueError("Length of strides my be 4.")
if tf_strides[0] != 1:
raise NotImplementedError('Strides on batch axis (N) must be 1.')
if tf_strides[3] != 1:
raise NotImplementedError('Strides on channel axis (C) must be 1.')
str_h, str_w = tf_strides[1], tf_strides[2]
# padding params
padding = tf_node.attr['padding'].s.decode("ascii")
pad_t, pad_b, pad_l, pad_r = tf_conv2d_pool_padding(
image.axes.lengths, weight.axes.lengths, tf_strides, padding)
if pad_t != pad_b or pad_l != pad_r:
raise NotImplementedError("Requires symmetric padding in ngraph:"
"pad_t(%s) == pad_b(%s) and"
"pad_l(%s) == pad_r(%s)" %
(pad_t, pad_b, pad_l, pad_r))
# conv params
params = dict(pad_d=0,
pad_h=pad_t,
pad_w=pad_l,
str_d=1,
str_h=str_h,
str_w=str_w)
# i, f, o axes
ax_i = ng.make_axes([ax_C, ax_D, ax_H, ax_W, ax_N])
ax_f = ng.make_axes([ax_C, ax_T, ax_R, ax_S, ax_K])
ax_o = ng.make_axes([
ng.make_axis(ax_K.length, name='C', roles=[ar.Channel]),
spatial_axis(ax_i, ax_f, params['pad_d'], params['str_d'],
ar.Depth),
spatial_axis(ax_i, ax_f, params['pad_h'], params['str_h'],
ar.Height),
spatial_axis(ax_i, ax_f, params['pad_w'], params['str_w'],
ar.Width), ax_N
])
# broadcast input / filter axes
image = ng.cast_axes(image, ng.make_axes([ax_N, ax_H, ax_W, ax_C]))
image = ng.expand_dims(image, ax_D, 1) # NHWC -> NDHWC
image = ng.axes_with_order(image, axes=ax_i) # NDHWC -> CDHWN
weight = ng.cast_axes(weight, ng.make_axes([ax_R, ax_S, ax_C, ax_K]))
weight = ng.expand_dims(weight, ax_T, 0) # RSCK -> TRSCK
weight = ng.axes_with_order(weight, axes=ax_f) # TRSCK -> CTRSK
# convolution
output = ng.convolution(params, image, weight, axes=ax_o)
# cast back to NHWC
oC, oD, oH, oW, oN = output.axes
output = ng.broadcast(output, ng.make_axes([oN, oD, oH, oW, oC]))
# slice away the oD
out_slicing = [slice(None), 0, slice(None), slice(None), slice(None)]
output = ng.Slice(output, out_slicing)
return output
def Conv(self, c2_op, inputs):
"""
Computes a 2-D convolution given 4D input and filter tensors.
Arguments:
c2_op: NodeDef object, the caffe2 node to convert.
inputs: List of ngraph Ops as inputs to this node.
Returns:
A ngraph Op corresponding to the caffe2 node.
Inputs to c2_op:
input, wegiths, filter
Supports caffe2's layout NHWC and NCHW as well.
"""
X, W, bias = inputs
order = [val.s for val in c2_op.arg if val.name == "order"]
if 1 != len(order):
raise ValueError("Multiple order values in convolution")
order = order[0]
if order not in ("NHWC", "NCHW"):
raise NotImplementedError("Unsupported order in convolution: {}",
order)
# set input axes shape
ax_N = ng.make_axis(name='N')
ax_C = ng.make_axis()
ax_D = ng.make_axis(length=1)
ax_H = ng.make_axis()
ax_W = ng.make_axis()
# set kernel axes shape
ax_kernel_D = ng.make_axis(length=1)
ax_kernel_H = ng.make_axis()
ax_kernel_W = ng.make_axis()
ax_kernel_ofm = ng.make_axis()
# create placeholders for output axes
oC = ng.make_axis(name='C')
oD = ng.make_axis(name='D', length=1)
oH = ng.make_axis(name='H')
oW = ng.make_axis(name='W')
axes_order = {
'NCHW': {
'X': [ax_N, ax_C, ax_H, ax_W],
'W': [ax_kernel_ofm, ax_C, ax_kernel_H, ax_kernel_W]
},
'NHWC': {
'X': [ax_N, ax_H, ax_W, ax_C],
'W': [ax_kernel_ofm, ax_kernel_H, ax_kernel_W, ax_C]
},
}
ng.make_axes(axes_order[order]['X']).set_shape(X.axes.lengths)
ng.make_axes(axes_order[order]['W']).set_shape(W.axes.lengths)
if 1 != len(bias.axes):
raise ValueError("Bias's must be 1D.")
if ax_kernel_ofm.length != bias.axes.lengths[0]:
raise ValueError(
"Bias's length must equal to number of output feature maps.")
# strides params
stride_size = [int(val.i) for val in c2_op.arg if val.name == "stride"]
if len(stride_size) != 1:
raise ValueError("Stride size must be scalar value")
str_h = str_w = stride_size[0]
# padding params
pad_t, pad_b, pad_l, pad_r = \
_c2_padding(c2_op,
in_NHWC=[ax_N.length, ax_H.length, ax_W.length, ax_C.length],
kernel_HWIO=[ax_kernel_H.length, ax_kernel_W.length,
ax_C.length, ax_kernel_ofm.length],
stride_NHWC=[1, str_h, str_w, 1])
if pad_t != pad_b or pad_l != pad_r:
raise NotImplementedError("Requires symmetric padding in ngraph:"
"pad_t(%s) == pad_b(%s) and"
"pad_l(%s) == pad_r(%s)" %
(pad_t, pad_b, pad_l, pad_r))
# conv params
params = dict(pad_d=0,
pad_h=pad_t,
pad_w=pad_l,
str_d=1,
str_h=str_h,
str_w=str_w,
dil_d=1,
dil_h=1,
dil_w=1)
# input, weight, output axes
internal_ax_dict = {
'X':
ng.make_axes([ax_C, ax_D, ax_H, ax_W, ax_N]),
'W':
ng.make_axes(
[ax_C, ax_kernel_D, ax_kernel_H, ax_kernel_W, ax_kernel_ofm])
}
oC.length = ax_kernel_ofm.length
oH.length = output_dim(ax_H.length, ax_kernel_H.length,
params['pad_h'], params['str_h'])
oW.length = output_dim(ax_W.length, ax_kernel_W.length,
params['pad_w'], params['str_w'])
internal_ax_dict['Y'] = ng.make_axes([oC, oD, oH, oW, ax_N])
# broadcast input / filter axes
# flow for NHWC order: | flow for NCHW order:
# input: | input:
# expand dims: NHWC -> NDHWC | expand dims: NCHW -> NDCHW
# reorder: NDHWC -> CDHWN | reorder: NDCHW -> CDHWN
# weights: | weights:
# expand dims: (ofm)HWC -> D(ofm)HWC | expand dims: (ofm)CHWC -> D(ofm)CHW
# reorder: D(ofm)HWC -> CDHW(ofm) | reorder: D(ofm)CHW -> CDHW(ofm)
X = ng.cast_axes(X, ng.make_axes(axes_order[order]['X']))
X = ng.expand_dims(X, ax_D, 1)
X = ng.axes_with_order(X, axes=internal_ax_dict['X'])
W = ng.cast_axes(W, ng.make_axes(axes_order[order]['W']))
W = ng.expand_dims(W, ax_kernel_D, 0)
W = ng.axes_with_order(W, axes=internal_ax_dict['W'])
# convolution
Y = ng.convolution(params, X, W, axes=internal_ax_dict['Y'])
# cast back to proper format
Y = ng.broadcast(Y, ng.make_axes([ax_N, oD, oH, oW, oC])) if "NHWC" == order \
else ng.broadcast(Y, ng.make_axes([ax_N, oD, oC, oH, oW])) # NCHW
# slice away the oD
out_slicing = [slice(None), 0, slice(None), slice(None), slice(None)]
Y = ng.tensor_slice(Y, out_slicing)
def _conv_bias_add(c2_op, inputs):
X, bias = inputs
bias = ng.cast_axes(bias,
axes=ng.make_axes(
[X.axes[1 if 'NCHW' == order else 3]]))
Y = ng.Add(X, bias)
return Y
return _conv_bias_add(c2_op, [Y, bias])
def test_conv_flatten_deriv(transformer_factory):
"""
Test deriv of conv followed by flatten
"""
# set shape
# NOTE: N must be >= 4 for GPU, but for CPU this could be decreased to
# speed up the test
N = 4
C, D, H, W = (3, 1, 28, 28)
T, R, S, K = (1, 5, 5, 8)
params = dict(pad_d=0,
pad_h=0,
pad_w=0,
str_d=1,
str_h=1,
str_w=1,
dil_d=1,
dil_h=1,
dil_w=1)
# i, f, o axes
ax_i = ng.make_axes([ax.C, ax.D, ax.H, ax.W, ax.N])
ax_f = ng.make_axes([ax.C, ax.T, ax.R, ax.S, ax.K])
ax_o = ng.make_axes([
ng.make_axis(roles=[ar.features_input]).named('C'),
ng.make_axis(roles=[ar.features_0]).named('D'),
ng.make_axis(roles=[ar.features_1]).named('H'),
ng.make_axis(roles=[ar.features_2]).named('W'), ax.N
])
ax_i.set_shape((C, D, H, W, N))
ax_f.set_shape((C, T, R, S, K))
ax_o.set_shape((K, D - T + 1, H - R + 1, W - S + 1, N))
axes_rsck = ng.make_axes([ax.R, ax.S, ax.C, ax.K])
axes_rsck_prime = ng.make_axes([
ng.make_axis(axis.length).named(axis.name + 'p') for axis in axes_rsck
])
axes_nmpqk = ng.make_axes([ax_o[-1], ax_o[1], ax_o[2], ax_o[3], ax_o[0]])
# broadcast input / filter axes
input_var = ng.variable(ax_i).named('input')
input_var.input = True
input_val = np.ones(input_var.axes.lengths)
filter_rsck_prime = ng.variable(axes_rsck_prime)
filter_var = filter_rsck_prime
filter_rsck = ng.cast_axes(filter_rsck_prime, axes_rsck)
filter_trsck = ng.expand_dims(filter_rsck, ax.T, 0)
filter_ctrsk = ng.axes_with_order(filter_trsck, axes=ax_f)
# convolution
output_kmpqn = ng.convolution(params, input_var, filter_ctrsk, axes=ax_o)
output_nmpqk = ng.axes_with_order(output_kmpqn, axes=axes_nmpqk)
# slice away the oD
out_slicing = [slice(None), 0, slice(None), slice(None), slice(None)]
output_npqk = ng.tensor_slice(output_nmpqk, out_slicing)
output = ng.flatten_at(output_npqk, idx=1)
# cost and grad
cost = ng.sum(output, out_axes=())
filter_var.input = True
filter_var.named('filter')
filter_val = np.ones(filter_var.axes.lengths)
with ExecutorFactory() as factory:
conv_comp = factory.executor(output, filter_var, input_var)
grad_filter_num_comp = factory.numeric_derivative(
cost, filter_var, 1.0, input_var)
grad_filter_sym_comp = factory.derivative(cost, filter_var, input_var)
grad_input_num_comp = factory.numeric_derivative(
cost, input_var, 1.0, filter_var)
grad_input_sym_comp = factory.derivative(cost, input_var, filter_var)
conv_val = conv_comp(filter_val, input_val)
conv_val_num = np.empty_like(conv_val)
conv_val_num.fill(C * T * R * S)
assert ng.testing.allclose(conv_val, conv_val_num)
grad_filter_num_val = grad_filter_num_comp(filter_val, input_val)
grad_filter_sym_val = grad_filter_sym_comp(filter_val, input_val)
assert ng.testing.allclose(grad_filter_num_val, grad_filter_sym_val)
grad_input_num_val = grad_input_num_comp(input_val, filter_val)
grad_input_sym_val = grad_input_sym_comp(input_val, filter_val)
assert ng.testing.allclose(grad_input_num_val, grad_input_sym_val)
def cifar_mean_subtract(x):
bgr_mean = ng.persistent_tensor(axes=x.axes[0],
initial_value=np.array([[104., 119.,
127.]]))
y = ng.expand_dims((x - bgr_mean) / 255., ax.D, 1)
return y
def Pool(self, c2_op, inputs):
"""
Performs max or average pooling on the input.
Arguments:
c2_op: NodeDef object, the tensorflow node to convert.
inputs: List of ngraph Ops as inputs to this node.
Returns:
A ngraph Op corresponding to the c2_op node.
Inputs to c2_op:
input
"""
supported_pooling = {'MaxPool': 'max', 'AveragePool': 'avg'}
image = inputs[0]
# TODO: we assume NCHW, make some assert here?
# set input axes shape
ax_N = ng.make_axis(name='N')
ax_C = ng.make_axis()
ax_D = ng.make_axis(length=1)
ax_H = ng.make_axis()
ax_W = ng.make_axis()
ng.make_axes([ax_N, ax_C, ax_H, ax_W]).set_shape(image.axes.lengths)
# create placeholders for output axes
oC = ng.make_axis(name='C')
oD = ng.make_axis(length=1, name='D')
oH = ng.make_axis(name='H')
oW = ng.make_axis(name='W')
# spatial kernel size
kernel_size = [int(val.i) for val in c2_op.arg if val.name == "kernel"]
if len(kernel_size) != 1:
raise ValueError("Kernel size must be scalar value")
# kernel is square
kernel_h = kernel_w = kernel_size[0]
kernel_d = kernel_c = 1
# strides params
stride_size = [int(val.i) for val in c2_op.arg if val.name == "stride"]
if len(stride_size) != 1:
raise ValueError("Stride size must be scalar value")
stride_h = stride_w = stride_size[0]
# padding params
pad_t, pad_b, pad_l, pad_r = \
_c2_padding(c2_op,
in_NHWC=[ax_N.length, ax_H.length, ax_W.length, ax_C.length],
kernel_HWIO=[kernel_h, kernel_w, ax_C.length, ax_C.length],
stride_NHWC=[1, stride_h, stride_w, 1])
if pad_t != pad_b or pad_l != pad_r:
raise NotImplementedError("Requires symmetric padding in ngraph:"
"pad_t(%s) == pad_b(%s) and"
"pad_l(%s) == pad_r(%s)" %
(pad_t, pad_b, pad_l, pad_r))
# pooling params
params = dict(op=supported_pooling[c2_op.type],
pad_d=0,
pad_h=pad_t,
pad_w=pad_l,
pad_c=0,
str_d=1,
str_h=stride_h,
str_w=stride_w,
str_c=1,
J=kernel_c,
T=kernel_d,
R=kernel_h,
S=kernel_w)
# i, o axes
oC.length = output_dim(ax_C.length, kernel_c, params['pad_c'],
params['str_c'])
oD.length = output_dim(ax_D.length, kernel_d, params['pad_d'],
params['str_d'])
oH.length = output_dim(ax_H.length, kernel_h, params['pad_h'],
params['str_h'])
oW.length = output_dim(ax_W.length, kernel_w, params['pad_w'],
params['str_w'])
ax_i = ng.make_axes([ax_C, ax_D, ax_H, ax_W, ax_N])
ax_o = ng.make_axes([oC, oD, oH, oW, ax_N])
# broadcast input / filter axes
image = ng.cast_axes(image, ng.make_axes([ax_N, ax_C, ax_H, ax_W]))
image = ng.expand_dims(image, ax_D, 1) # NCHW -> NDCHW
image = ng.axes_with_order(image, axes=ax_i) # NDCHW -> CDHWN
# pooling
output = ng.pooling(params, image, axes=ax_o)
# cast back to NDCHW
output = ng.broadcast(output, ng.make_axes([ax_N, oD, oC, oH, oW]))
# slice away the oD
out_slicing = [slice(None), 0, slice(None), slice(None), slice(None)]
output = ng.tensor_slice(output, out_slicing)
return output
