def padding(input_tensor, padding_size, name="padding"):
"""Construct a tensor by padding a given tensor.
Args:
input_tensor: Input tensor.
padding_size: A list in the format of {dim0_begin, dim0_end, dim1_begin,
dim1_end, ...} that represent number of values padded to
each dimension. Note that the order of dimensions of this
must align with the data layout of input_tensor.
name: Name of the operator.
Returns:
A padded version of the input tensor.
"""
src_layout = input_tensor.shape.layout
src_dims = input_tensor.shape.dims
if len(padding_size) != 2 * len(src_dims):
raise ValueError(
"len(padding_size) should be 2x input_tensor.shape.dims")
output_tensor_dims = [0] * len(src_dims)
for i in range(len(src_dims)):
output_tensor_dims[i] = src_dims[i] + padding_size[
2 * i] + padding_size[2 * i + 1]
params = node_pb2.Params()
params.padding_params.padding_size.extend(padding_size)
return common.add_node(name=name,
op=types_pb2.Padding,
input_tensors=[input_tensor],
output_tensors_dims=[output_tensor_dims],
output_tensor_layout=input_tensor.shape.layout,
params=params)[0]
def concat(input_tensors, axis=0, name="concat"):
"""Concatenate tensors into one.
Args:
input_tensors: Input tensor to be concatenated.
axis: The dimension along which to concatenate.
name: Name of the operator.
Returns:
A concatenated tensor.
"""
dims = np.delete(input_tensors[0].shape.dims, axis)
if not all([
np.array_equal(np.delete(x.shape.dims, axis), dims)
for x in input_tensors
]):
raise ValueError(
"Tensors must have the same shape, except in axis %d along which to "
"concatenate." % axis)
output_tensor_dims = list(input_tensors[0].shape.dims)
output_tensor_dims[axis] = sum(x.shape.dims[axis] for x in input_tensors)
params = node_pb2.Params()
params.concat_params.concat_axis = axis
return common.add_node(name=name,
op=types_pb2.Concat,
input_tensors=input_tensors,
output_tensors_dims=[output_tensor_dims],
output_tensor_layout=input_tensors[0].shape.layout,
params=params)[0]
def mat_mul(
input_tensor, weight_tensor, activation=None, activation_params=None,
name="mat_mul"):
"""Compute a matrix multiplication for `input_tensor` and `weight_tensor`.
Args:
input_tensor: A 2D `Tensor`. Shaped as `NC`, where `N` is batch size and `C`
is number of channels.
weight_tensor: A 2D `Tensor`. Shaped as `NC` or `CN`, where `N` is number of
neurons and `C` is the same as in `input_tensor`.
activation/activation_params: Activation function to use (optional).
name: Operator name (optional).
"""
input_tensor, weight_tensor = array_ops.check_and_add_layout_transform(
name=name, op=types_pb2.InnerProduct,
input_tensors=[input_tensor, weight_tensor])
weight_layout = weight_tensor.shape.layout
actIdx = 1 if weight_layout == types_pb2.NC else 0
neuronIdx = 0 if weight_layout == types_pb2.NC else 1
assert (len(input_tensor.shape.dims) == 2
and len(weight_tensor.shape.dims) == 2
and input_tensor.shape.dims[1] == weight_tensor.shape.dims[actIdx])
output_tensor_dims = [
input_tensor.shape.dims[0], weight_tensor.shape.dims[neuronIdx]
]
params = node_pb2.Params()
if activation is not None:
params.act_params.CopyFrom(
activation_ops.to_proto(activation, activation_params))
return common.add_node(
name=name, op=types_pb2.InnerProduct,
input_tensors=[input_tensor, weight_tensor],
output_tensors_dims=[output_tensor_dims],
output_tensor_layout=types_pb2.NC, params=params)[0]
def relu(input_tensor, name="relu"):
"""Rectified linear unit operator."""
return common.add_node(name=name,
op=types_pb2.ReLU,
input_tensors=[input_tensor],
output_tensors_dims=[input_tensor.shape.dims],
output_tensor_layout=input_tensor.shape.layout)[0]
def tanh(input_tensor, name="tanh"):
"""Tanh operator."""
return common.add_node(name=name,
op=types_pb2.Tanh,
input_tensors=[input_tensor],
output_tensors_dims=[input_tensor.shape.dims],
output_tensor_layout=input_tensor.shape.layout)[0]
def convolution(
input_tensor, filter_tensor, stride, padding, activation=None,
activation_params=None, name="conv"):
"""Compute a 3D Convolution given 4D `input_tensor` and `filter_tensor`.
Args:
input_tensor: A 4D `Tensor`.
filter_tensor: A 4D `Tensor`.
stride: A list of two integers: [row_stride, col_stride].
padding: A string from: `same`, `valid`. The zero padding options.
activation: A string representing the activation function (optional).
activation_params: kwargs for the activation function (optional).
name: Operator name (optional).
"""
def compute_output_dim(input_dim, weight_dim, stride, padding):
pad = 0
if to_padding_type(padding) == types_pb2.SamePadding:
pad = weight_dim - 1
return (input_dim - weight_dim + pad) // stride + 1
input_tensor, filter_tensor = array_ops.check_and_add_layout_transform(
name=name, op=types_pb2.Convolution3d,
input_tensors=[input_tensor, filter_tensor])
row_idx = 2 if input_tensor.shape.layout == types_pb2.NCHW else 1
col_idx = 3 if input_tensor.shape.layout == types_pb2.NCHW else 2
chan_idx = 1 if input_tensor.shape.layout == types_pb2.NCHW else 3
assert input_tensor.dims(chan_idx) == filter_tensor.dims(chan_idx), (
"The weights must have the same number of channels as the inputs.")
output_rows = compute_output_dim(input_tensor.shape.dims[row_idx],
filter_tensor.shape.dims[row_idx], stride[0],
padding)
output_cols = compute_output_dim(input_tensor.shape.dims[col_idx],
filter_tensor.shape.dims[col_idx], stride[1],
padding)
output_layout = input_tensor.shape.layout
if output_layout == types_pb2.NCHW:
output_tensor_dims = [
input_tensor.shape.dims[0], filter_tensor.shape.dims[0], output_rows,
output_cols
]
elif output_layout == types_pb2.NHWC:
output_tensor_dims = [
input_tensor.shape.dims[0], output_rows, output_cols,
filter_tensor.shape.dims[0]
]
else:
assert False, "Unsupported output layout!"
params = node_pb2.Params()
params.conv_params.padding = to_padding_type(padding)
params.conv_params.stride.extend(stride)
if activation is not None:
params.act_params.CopyFrom(
activation_ops.to_proto(activation, activation_params))
return common.add_node(
name=name, op=types_pb2.Convolution3d,
input_tensors=[input_tensor, filter_tensor],
output_tensors_dims=[output_tensor_dims],
output_tensor_layout=output_layout, params=params)[0]
def softmax(input_tensor, name=None):
"""Softmax operator."""
input_tensor = array_ops.check_and_add_layout_transform(
name=name, op=types_pb2.Softmax, input_tensors=[input_tensor])[0]
return common.add_node(
name=name, op=types_pb2.Softmax, input_tensors=[input_tensor],
output_tensors_dims=[input_tensor.shape.dims],
output_tensor_layout=input_tensor.shape.layout)[0]
def lrelu(input_tensor, slope=0.2, name="lrelu"):
"""Leaky rectified linear unit operator: max(slope * x, 0)."""
params = node_pb2.Params()
params.act_params.lrelu_params.slope = slope
return common.add_node(
name=name, op=types_pb2.LReLU, input_tensors=[input_tensor],
output_tensors_dims=[input_tensor.shape.dims],
output_tensor_layout=input_tensor.shape.layout, params=params)[0]
def sigmoid(input_tensor, name="sigmoid"):
"""Sigmoid operator.
Defined as 1/(1 + exp(-input_tensor)).
"""
return common.add_node(
name=name, op=types_pb2.Sigmoid, input_tensors=[input_tensor],
output_tensors_dims=[input_tensor.shape.dims],
output_tensor_layout=input_tensor.shape.layout)[0]
def _math_op_common(tensor_a, tensor_b, op, name, output_tensor_dtype=None):
if tensor_a.shape.dims != tensor_b.shape.dims:
tensor_a, tensor_b = common.broadcast_inputs(tensor_a, tensor_b, name)
if output_tensor_dtype == None:
output_tensor_dtype = tensor_a.data_type
return common.add_node(
name=name, op=op, input_tensors=[tensor_a, tensor_b],
output_tensors_dims=[tensor_a.shape.dims],
output_tensor_layout=tensor_a.shape.layout,
output_tensor_dtype=output_tensor_dtype)[0]
def input_data(input_tensor, name="data"):
"""Create an data node for the given tensor.
This effectively turns a `Tensor` into a `Node`.
"""
return common.add_node(name=name,
op=types_pb2.Data,
input_tensors=[input_tensor],
output_tensors_dims=[input_tensor.shape.dims],
output_tensor_layout=input_tensor.shape.layout)[0]
def selu(input_tensor, alpha=1.6733, lambda_param=1.0507, name="selu"):
"""Scaled exponential linear unit function.
Defined as: lambda_param * elu(input_tensor, alpha).
"""
params = node_pb2.Params()
params.act_params.elu_params.alpha = alpha
params.act_params.elu_params.lambda_param = lambda_param
return common.add_node(
name=name, op=types_pb2.SELU, input_tensors=[input_tensor],
output_tensors_dims=[input_tensor.shape.dims],
output_tensor_layout=input_tensor.shape.layout, params=params)[0]
def hard_tanh(input_tensor, min=-1, max=1, name="hard_tanh"):
"""Hard tanh operator.
This bounds the min and max values of the tanh operator.
"""
params = node_pb2.Params()
params.act_params.hard_tanh_params.min = min
params.act_params.hard_tanh_params.max = max
return common.add_node(
name=name, op=types_pb2.HardTanh, input_tensors=[input_tensor],
output_tensors_dims=[input_tensor.shape.dims],
output_tensor_layout=input_tensor.shape.layout, params=params)[0]
def elu(input_tensor, alpha=0.1, name="relu"):
"""Exponential linear unit function.
Defined as:
if input_tensor > 0, alpha * exp(input_tensor - 1), else input_tensor.
"""
params = node_pb2.Params()
params.act_params.elu_params.alpha = alpha
return common.add_node(
name=name, op=types_pb2.ELU, input_tensors=[input_tensor],
output_tensors_dims=[input_tensor.shape.dims],
output_tensor_layout=input_tensor.shape.layout, params=params)[0]
def merge(input_tensors, name="merge"):
"""Forward the value of an available tensor from inputs to output.
Args:
input_tensors: Input tensors. All are dead tensor except one.
Returns:
A tensor that the available input tensor forwards to.
"""
return common.add_node(
name=name,
op=types_pb2.Merge,
input_tensors=input_tensors,
output_tensors_dims=[input_tensors[0].shape.dims],
output_tensor_layout=input_tensors[0].shape.layout)[0]
def reorder(input_tensor, target_layout, name="reorder"):
"""Reorder the data of a given `Tensor` with the target layout.
Args:
input_tensor: A `Tensor`.
target_layout: The target layout.
name: Operator name (optional).
Returns:
A new `Tensor` with the layout as `target_layout`.
"""
src_layout = input_tensor.shape.layout
src_dims = input_tensor.shape.dims
if src_layout == types_pb2.NCHW:
assert (target_layout == types_pb2.NHWC
or target_layout == types_pb2.NC)
if target_layout == types_pb2.NC:
output_tensor_dims = [src_dims[0], np.prod(src_dims[1:])]
else:
output_tensor_dims = [
src_dims[0], src_dims[2], src_dims[3], src_dims[1]
]
elif src_layout == types_pb2.NHWC:
assert (target_layout == types_pb2.NCHW
or target_layout == types_pb2.NC)
if target_layout == types_pb2.NC:
output_tensor_dims = [src_dims[0], np.prod(src_dims[1:])]
else:
output_tensor_dims = [
src_dims[0], src_dims[3], src_dims[1], src_dims[2]
]
elif (src_layout == types_pb2.NTC and target_layout == types_pb2.NCT) or (
src_layout == types_pb2.NCT and target_layout == types_pb2.NTC):
output_tensor_dims = [src_dims[0], src_dims[2], src_dims[1]]
elif (src_layout == types_pb2.NC and target_layout == types_pb2.CN) or (
src_layout == types_pb2.CN and target_layout == types_pb2.NC):
# 2D tensor transposition.
output_tensor_dims = [src_dims[1], src_dims[0]]
else:
raise ValueError(
"Unsupported reordering %s->%s!" %
(DataLayout.Name(src_layout), DataLayout.Name(target_layout)))
return common.add_node(name=name,
op=types_pb2.Reorder,
input_tensors=[input_tensor],
output_tensors_dims=[output_tensor_dims],
output_tensor_layout=target_layout)[0]
def split(input_tensor, num_or_size_splits, axis=0, name="split"):
"""Split a tensor into sub tensors.
Args:
input_tensor: Input tensor.
num_or_size_splits: Either an integer indicating the number of splits along
axis or a 1D list containing the sizes of each output tensor along axis.
If an integer, then it must evenly divide input_tensor.shape.dims[axis];
otherwise the sum of sizes along the split axis must match that of the
value.
axis: The dimension to split.
name: Name of the operator.
Returns:
A list of sub tensors.
"""
splits = num_or_size_splits
dim = input_tensor.shape.dims[axis]
if not isinstance(num_or_size_splits, list):
if dim % num_or_size_splits != 0:
raise ValueError(
"The size (%d) of the axis along which to split must divide the "
"splits (%d)!" % (dim, num_or_size_splits))
splits = [dim // num_or_size_splits] * num_or_size_splits
if sum(splits) != input_tensor.shape.dims[axis]:
raise ValueError(
"the sum (%d) of sizes along the split axis must match that of the "
"input (%d)!" % (sum(splits), input_tensor.shape.dims[axis]))
if splits == [1]:
warnings.warn(
"Number of splits is 1 for the split operator, thus this operator is "
"optimized out.")
return [input_tensor]
output_tensors_dims = []
for s in splits:
dims = list(input_tensor.shape.dims)
dims[axis] = s
output_tensors_dims.append(dims)
params = node_pb2.Params()
params.split_params.split_axis = axis
return common.add_node(name=name,
op=types_pb2.Split,
input_tensors=[input_tensor],
output_tensors_dims=output_tensors_dims,
output_tensor_layout=input_tensor.shape.layout,
params=params)
def switch(input_tensor, pred, name="switch"):
"""Forward the input to output port determined by the given predication.
Args:
input_tensor: Input tensor.
pred: Predication tensor. The tensor should only contain a single boolean
value.
Returns:
output_false, output_true: Two tensors representing the two branches of the
switch. Input will only be forwarded to the taken branch.
"""
return common.add_node(name=name,
op=types_pb2.Switch,
input_tensors=[input_tensor, pred],
output_tensors_dims=[input_tensor.shape.dims] * 2,
output_tensor_layout=input_tensor.shape.layout)
def reshape(input_tensor, shape, layout, name="reshape"):
""" Reshape the given tensor in the same order.
Args:
input_tensor: Input tensor.
shape: New shape.
layout: New layout.
name: Name of the operator.
Returns:
Tensor with the new shape.
"""
assert np.prod(input_tensor.shape.dims) == np.prod(shape)
return common.add_node(name=name,
op=types_pb2.Reshape,
input_tensors=[input_tensor],
output_tensors_dims=[shape],
output_tensor_layout=layout)[0]
def batch_norm(
input_tensor, mean_tensor, var_tensor, gamma_tensor, beta_tensor,
activation=None, activation_params=None, name="batch_norm"):
"""Perform batch normalization.
Args:
input_tensor: A 2D or 4D `Tensor`.
mean_tensor: Mean parameter.
var_tensor: Variance parameter. For performance reasons, this is
precomputed as 1/sqrt(variance + eps).
gamma_tensor: Gamma parameter.
beta_tensor: Beta parameter.
activation/activation_params: Activation function to use (optional).
name: Operator name (optional).
"""
assert (len(mean_tensor.shape.dims) == 2 and len(var_tensor.shape.dims) == 2
and len(gamma_tensor.shape.dims) == 2
and len(beta_tensor.shape.dims) == 2)
# If the batch norm is after a FC layer, then the input/output tensors should
# be in NC. Otherwise, the batch norm is after a convolution layer, and we
# check backend_layouts for expected input/output layouts and do layout
# transformation if needed.
post_fc = False
if len(input_tensor.shape.dims) == 2:
post_fc = True
if not post_fc:
input_tensor = array_ops.check_and_add_layout_transform(
name=name, op=types_pb2.BatchNorm, input_tensors=[input_tensor])[0]
output_layout = types_pb2.UnknownLayout
output_layout = types_pb2.NC if post_fc else input_tensor.shape.layout
params = node_pb2.Params()
if activation is not None:
params.act_params.CopyFrom(
activation_ops.to_proto(activation, activation_params))
return common.add_node(
name=name, op=types_pb2.BatchNorm, input_tensors=[
input_tensor, mean_tensor, var_tensor, gamma_tensor, beta_tensor
], output_tensors_dims=[input_tensor.shape.dims],
output_tensor_layout=output_layout, params=params)[0]
def max_pool(input_tensor, pool_size, stride, name="max_pool"):
"""Compute max pooling.
Args:
input_tensor: A 4D `Tensor`.
pool_size: A list of two integers: [pool_rows, pool_cols].
stride: A list of two integers: [row_stride, col_stride].
name: Operator name (optional).
"""
def compute_output_dim(input_dim, pool_size, stride):
return (input_dim - pool_size) // stride + 1
input_tensor = array_ops.check_and_add_layout_transform(
name=name, op=types_pb2.MaxPooling, input_tensors=[input_tensor])[0]
row_idx = 2 if input_tensor.shape.layout == types_pb2.NCHW else 1
col_idx = 3 if input_tensor.shape.layout == types_pb2.NCHW else 2
output_rows = compute_output_dim(input_tensor.shape.dims[row_idx],
pool_size[0], stride[0])
output_cols = compute_output_dim(input_tensor.shape.dims[col_idx],
pool_size[1], stride[1])
output_layout = input_tensor.shape.layout
if output_layout == types_pb2.NCHW:
output_tensor_dims = [
input_tensor.shape.dims[0], input_tensor.shape.dims[1], output_rows,
output_cols
]
else:
output_tensor_dims = [
input_tensor.shape.dims[0], output_rows, output_cols,
input_tensor.shape.dims[3]
]
params = node_pb2.Params()
params.pool_params.stride.extend(stride)
params.pool_params.pool_size.extend(pool_size)
return common.add_node(
name=name, op=types_pb2.MaxPooling, input_tensors=[input_tensor],
output_tensors_dims=[output_tensor_dims],
output_tensor_layout=output_layout, params=params)[0]
def repeat(input_tensor, multiples, name="repeat"):
"""Construct a tensor by repeating a given tensor.
Args:
input_tensor: Input tensor.
multiples: A list that represents the number of multiples in each dimension
of the input tensor.
name: Name of the operator.
Returns:
A repeated version of the input tensor.
"""
if len(input_tensor.shape.dims) != len(multiples):
raise ValueError(
"The multiples of the repeat operator must have the same number of "
"dims as the input tensor.")
output_tensor_dims = np.multiply(input_tensor.shape.dims, multiples)
return common.add_node(name=name,
op=types_pb2.Repeat,
input_tensors=[input_tensor],
output_tensors_dims=[output_tensor_dims],
output_tensor_layout=input_tensor.shape.layout)[0]
