zhutian
264 lines
10 KiB
Python
264 lines
10 KiB
Python
from __future__ import division
|
|
|
|
from collections import namedtuple
|
|
import numpy as np
|
|
import tensorflow as tf
|
|
|
|
import itertools
|
|
|
|
|
|
pad_ops = namedtuple("pad_ops",
|
|
["max_op", "ceil_op", "floor_op", "cast_int_op"])
|
|
|
|
pad_numpy_ops = pad_ops(np.maximum, np.ceil, np.floor,
|
|
lambda arr: arr.astype(np.int64))
|
|
pad_tf_ops = pad_ops(tf.maximum, tf.math.ceil, tf.math.floor,
|
|
lambda tensor: tf.cast(tensor, tf.int64))
|
|
|
|
|
|
def calc_pads_same(in_spatial_shape, kernel_shape, strides,
|
|
dilations, padding, padding_ops=pad_numpy_ops,
|
|
pads_order=1):
|
|
"""
|
|
Calculates the SAME paddings that need to be added to the input
|
|
|
|
Args:
|
|
in_spatial_shape: input spatial shape
|
|
kernel_shape: the size of the kernel along each axis
|
|
strides: stride along each spatial axis
|
|
dilations: dilations value along each spatial axis
|
|
padding: padding to calculate: SAME_UPPER or
|
|
SAME_LOWER
|
|
padding_ops: namedtuple with ops to be used during
|
|
calculations. there are two sets of ops
|
|
defined pad_numpy_ops and pad_tf_ops with
|
|
numpy and tensorflow ops
|
|
pads_order: order of returned pads. possible options are:
|
|
1 - b1, b2, ..., bn, e1, e2, ..., en
|
|
2 - b1, e1, b2, e2, ..., bn, en
|
|
where n = len(kernel_shape) * 2,
|
|
b1, b2, ..., bn define pads at the begging of
|
|
axis
|
|
e1, e2, ..., en define pads at the end of
|
|
axis
|
|
Return:
|
|
pads: array with calculated pads. the order of the
|
|
values is determined by `pads_order`
|
|
|
|
"""
|
|
spatial_size = len(kernel_shape)
|
|
pads = [0] * (spatial_size * 2)
|
|
for i in range(spatial_size):
|
|
in_size = in_spatial_shape[i]
|
|
filter_size = (kernel_shape[i] - 1) * dilations[i] + 1
|
|
|
|
out_size = padding_ops.ceil_op(in_size / strides[i])
|
|
out_size = padding_ops.cast_int_op(out_size)
|
|
pad_along_axis = \
|
|
padding_ops.max_op((out_size - 1) * strides[i] +
|
|
filter_size - in_size, 0)
|
|
if padding.lower() == "same_lower":
|
|
pad_op = padding_ops.ceil_op
|
|
else:
|
|
pad_op = padding_ops.floor_op
|
|
pad_begin = pad_op(pad_along_axis / 2)
|
|
|
|
pad_begin = padding_ops.cast_int_op(pad_begin)
|
|
pad_along_axis = padding_ops.cast_int_op(pad_along_axis)
|
|
|
|
pad_end = pad_along_axis - pad_begin
|
|
|
|
pads[i * pads_order] = pad_begin
|
|
pads[i * pads_order +
|
|
(spatial_size if pads_order == 1 else 1)] = pad_end
|
|
|
|
return pads
|
|
|
|
|
|
def calc_output_shape(input_spatial_shape, kernel_shape, strides, dilations,
|
|
padding, ceil_mode=False):
|
|
"""
|
|
Calculate output shape
|
|
|
|
Args:
|
|
input_spatial_shape: input spatial shape
|
|
kernel_shape: the size of the kernel along each axis
|
|
strides: stride along each spatial axis
|
|
dilations: dilations value along each spatial axis
|
|
padding: can be explicit paddings, "SAME_UPPER" or
|
|
"SAME_LOWER"
|
|
Return:
|
|
output_shape: calculated output shape
|
|
"""
|
|
spatial_size = len(input_spatial_shape)
|
|
|
|
if type(padding) is not list and type(padding) is not np.ndarray:
|
|
if padding.lower().startswith("same"):
|
|
padding = calc_pads_same(input_spatial_shape, kernel_shape,
|
|
strides, dilations, padding)
|
|
else:
|
|
padding = [0] * spatial_size * 2
|
|
|
|
output_shape = []
|
|
for dim in range(spatial_size):
|
|
output_shape.append(_pooling_output_shape(input_spatial_shape[dim],
|
|
kernel_shape[dim], strides[dim], dilations[dim],
|
|
padding[dim] + padding[dim + spatial_size],
|
|
ceil_mode))
|
|
|
|
return output_shape
|
|
|
|
|
|
def _pooling_output_shape(input_size, ksize, stride, dilation, pad, ceil_mode):
|
|
output_size = (input_size + pad - ((ksize - 1) * dilation + 1) +
|
|
((stride-1) if ceil_mode else 0)) // stride + 1
|
|
if (pad):
|
|
if ((output_size - 1) * stride >= input_size + pad):
|
|
output_size -= 1
|
|
return output_size
|
|
|
|
|
|
def py_pool(input, kernel_shape, strides=None, dilations=None,
|
|
padding=None, ceil_mode=False, pooling_type="MAX",
|
|
include_indices=True, p=2):
|
|
"""
|
|
Implementation of Max and Average pool operations in Python
|
|
Args:
|
|
input: input N-D data array in NC* format
|
|
kernel_shape: the size of the kernel along each axis
|
|
strides: stride along each spatial axis
|
|
dilations: dilations value along each spatial axis of filter
|
|
padding: padding for the beginning and ending along each
|
|
spatial axis. `padding` format should be as follow
|
|
[x1_begin, x2_begin...x1_end, x2_end,...]
|
|
ceil_mode: whether to use ceil or floor (default) to compute
|
|
the output shape.
|
|
pooling_type: specifies pooling type. Values can be "MAX", "AVG" or
|
|
"LP"
|
|
include_indices: should indices be included in the output
|
|
p: specifies the p parameter for LpPooling
|
|
Return:
|
|
pooled: output data from max pooling across the input
|
|
ind: indices of the selected max values from the input
|
|
"""
|
|
|
|
if type(pooling_type) is not str:
|
|
pooling_type = pooling_type.decode("UTF-8")
|
|
|
|
input_shape = np.shape(input)
|
|
inp_sp_shape = input_shape[2:]
|
|
input_dtype = input.dtype
|
|
if np.issubdtype(input_dtype, np.integer):
|
|
input_dtype_min = np.iinfo(input_dtype).min
|
|
else:
|
|
input_dtype_min = np.finfo(input_dtype).min
|
|
|
|
if pooling_type == "LP":
|
|
rootN = (1.0 / p)
|
|
|
|
def _loop_over_output(batch, channel):
|
|
dims = [range(output_sp_shape[d]) for d in range(spatial_size)]
|
|
for counters in itertools.product(*dims):
|
|
input_ranges = []
|
|
for dim in range(spatial_size):
|
|
dim_start = \
|
|
counters[dim] * strides[dim] - pads[dim * 2]
|
|
dim_end = \
|
|
min(dim_start + (kernel_shape[dim] - 1) * dilations[dim]
|
|
+ 1, inp_sp_shape[dim])
|
|
while dim_start < 0:
|
|
dim_start += dilations[dim]
|
|
|
|
cur_range = [i for i in range(dim_start,
|
|
dim_end, dilations[dim])]
|
|
input_ranges.append(cur_range)
|
|
if pooling_type in ["AVG", "LP"]:
|
|
val_sum = 0
|
|
val_count = 0
|
|
else:
|
|
maxval = input_dtype_min
|
|
maxind = -1
|
|
for input_ind in itertools.product(*input_ranges):
|
|
ind = (batch, channel) + input_ind
|
|
val = input[ind]
|
|
if pooling_type == "AVG":
|
|
val_sum += val
|
|
val_count += 1
|
|
elif pooling_type == "LP":
|
|
val_sum += abs(val ** p)
|
|
else:
|
|
if val > maxval:
|
|
maxval = val
|
|
ind = 0
|
|
for i in range(spatial_size):
|
|
coef = 1
|
|
for j in range(i+1, spatial_size):
|
|
coef *= inp_sp_shape[j]
|
|
ind += input_ind[i] * coef
|
|
maxind = ind
|
|
ind = (batch, channel) + counters
|
|
if pooling_type == "AVG":
|
|
out_pool[ind] = val_sum / val_count
|
|
elif pooling_type == "LP":
|
|
out_pool[ind] = val_sum ** rootN
|
|
else:
|
|
out_pool[ind] = maxval
|
|
out_ind[ind] = maxind
|
|
|
|
spatial_size = len(kernel_shape)
|
|
|
|
batch_size = input_shape[0]
|
|
channels_num = input_shape[1]
|
|
|
|
if strides is None:
|
|
strides = kernel_shape
|
|
|
|
if dilations is None:
|
|
dilations = [1] * spatial_size
|
|
|
|
if padding is None:
|
|
padding = [0] * spatial_size * 2
|
|
|
|
if type(padding) is bytes:
|
|
padding = padding.decode()
|
|
|
|
if type(padding) is not list and type(padding) is not np.ndarray:
|
|
if type(padding) is not str:
|
|
padding = padding.decode("UTF-8")
|
|
|
|
if padding.lower().startswith("same"):
|
|
padding = calc_pads_same(inp_sp_shape, kernel_shape, strides,
|
|
dilations, padding)
|
|
else:
|
|
padding = [0] * spatial_size * 2
|
|
|
|
pads = []
|
|
pad_along_axis = []
|
|
output_sp_shape = []
|
|
|
|
for dim in range(spatial_size):
|
|
pads.append(padding[dim])
|
|
pads.append(padding[dim + spatial_size])
|
|
pad_along_axis.append(padding[dim] + padding[dim + spatial_size])
|
|
|
|
input_size = input_shape[dim + 2]
|
|
output_size = \
|
|
_pooling_output_shape(input_size, kernel_shape[dim],
|
|
strides[dim], dilations[dim],
|
|
pad_along_axis[dim], ceil_mode)
|
|
output_sp_shape.append(output_size)
|
|
|
|
out_pool = np.zeros([input_shape[0], input_shape[1]] +
|
|
output_sp_shape, input_dtype)
|
|
out_ind = np.zeros([input_shape[0], input_shape[1]] +
|
|
output_sp_shape, np.int64)
|
|
|
|
for batch in range(batch_size):
|
|
for channel in range(channels_num):
|
|
_loop_over_output(batch, channel)
|
|
|
|
if not include_indices:
|
|
return out_pool
|
|
else:
|
|
return out_pool, out_ind
|