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