add pt2tf tool
This commit is contained in:
zxros10
@@ -0,0 +1,702 @@
|
||||
from __future__ import division
|
||||
|
||||
import tensorflow as tf
|
||||
import numpy as np
|
||||
|
||||
from onnx_tf.common import pooling_helper
|
||||
from onnx_tf.common.tf_helper import tf_shape
|
||||
from onnx_tf.common.tf_helper import tf_product
|
||||
|
||||
|
||||
class DilatedPooling(object):
|
||||
"""
|
||||
This class implements two main methods:
|
||||
dilated_pool:
|
||||
calculates a max or average pool over the input
|
||||
|
||||
dilated_maxpool_with_argmax:
|
||||
calculates a maxpool over the input and returns the
|
||||
indices/argmax of the selected values
|
||||
|
||||
In addition to the standard features of pooling operations in
|
||||
Tensorflow, these methods support dilations, ceil mode, SAME_LOWER and
|
||||
explicit padding.
|
||||
|
||||
Dilations are partly supported in Tensorflow in `tf.nn.pool` and
|
||||
`tf.nn.dilation2d`. The code will try to use the Tensoflow build-in
|
||||
functions as much as poosible.
|
||||
|
||||
In cases, not supported by Tensorflow there is a custom algorith of
|
||||
dilated pooling `_remove_dilations`.
|
||||
|
||||
The idea behind `_remove_dilations` is to transform the input N-D data
|
||||
into a supported input for the standard tf.nn.pool operation.
|
||||
This is achieved by calculating N-D indicies for the values which will
|
||||
be selected from the input when applying the dilations and
|
||||
then extracting the values using tf.gather_nd. Next step is to execute
|
||||
`tf.nn.pool` on this new input data with **strides=kernel_shape** and
|
||||
no dilations. The resulting pool will be the result we are looking for.
|
||||
|
||||
In case of `deilated_maxpool_with_argmax` an additional step is needed
|
||||
to recalculated the resulting indices back into the original
|
||||
data indices. It is done with `_calc_orig_argmax`
|
||||
|
||||
Here is a simple example of how the algorithm works:
|
||||
|
||||
kernel_shape = [3]
|
||||
strides = [2]
|
||||
dilations = [3]
|
||||
|
||||
Input 1D data:
|
||||
|
||||
x-----x-----x-----x-----x-----x-----x-----x-----x-----x-----x
|
||||
| * | | ** | * | | ** | * | | ** | |
|
||||
| 10 | 9 | 30 | 7 | 6 | 15 | 16 | 17 | 18 | 19 |
|
||||
x-----x-----x-----x-----x-----x-----x-----x-----x-----x-----x
|
||||
(0) (1) (2) (3) (4) (5) (6) (7) (8) (9)
|
||||
|
||||
where * represents the values selected during the first sliding window
|
||||
step and ** during the second sliding window step
|
||||
|
||||
the resulting indices will be:
|
||||
|
||||
[0, 3, 6, 2, 5, 8]
|
||||
| | | |
|
||||
First Second
|
||||
step step
|
||||
|
||||
after tf.gather_nd operation we get a new input data with
|
||||
removed dilations:
|
||||
|
||||
[10, 7, 16, 30, 15, 18]
|
||||
|
||||
and apllying tf.nn.maxpool (or avgpool) with strides = kernel_shape = 3
|
||||
will result into:
|
||||
|
||||
[16, 30]
|
||||
|
||||
which is the result of the dilated maxpooling.
|
||||
|
||||
Here is pseudo code of the algorithm with comments:
|
||||
|
||||
FUNCTION _remove_dilations:
|
||||
/* Calculate N-D index of the values to be selected by the
|
||||
dilations and strides */
|
||||
|
||||
/* Do a loop over the input spatial dimensions starting from the
|
||||
last (most internal) going up to the first dimension
|
||||
|
||||
On every step of the loop calculate the input indices and
|
||||
"combine" them with the already calculated indices from the
|
||||
previous dimensions using cartesian product.
|
||||
*/
|
||||
LOOP with **dimension** from **dimensions_count** to **0**:
|
||||
|
||||
// Initialize empty gather_nd index
|
||||
gather_ind = []
|
||||
|
||||
// Calculate the output size for the current dimension
|
||||
dim_filter_size = (dim_kernel_size - 1) * dim_dilations
|
||||
dim_output_size = (((dim_input_size - dim_filter_size) //
|
||||
dim_strides) + 1) * dim_kernel_size)
|
||||
|
||||
/* For every output index, calculate the corresponding index
|
||||
into the input data */
|
||||
dim_input_indices = range(0, dim_output_size)
|
||||
dim_input_indices = calculate_input_indicies(dim_input_indices)
|
||||
|
||||
/* combine the calculated indices with the previous dimensions
|
||||
*/
|
||||
gather_ind = cartesian_product(dim_input_indices, gather_ind)
|
||||
END LOOP
|
||||
|
||||
/* For example for 2D input the resulting gather_ind will
|
||||
look like this:
|
||||
|
||||
[[y1, x1], [y2, x2], ..., [yn, xm]]
|
||||
|
||||
where:
|
||||
n is the height
|
||||
m is the width and
|
||||
[xi, yi] are the 2D indices in the input data
|
||||
*/
|
||||
|
||||
new_data = tf.gather_nd(input, gather_ind)
|
||||
|
||||
reshape new_data to the correct output shape
|
||||
|
||||
RETURN new_data
|
||||
|
||||
|
||||
Before executing _remove_dilations the code will apply paddings to the
|
||||
input data if needed. Padding is done using tf.pad with -inf values.
|
||||
Check `_remove_dilations` code for more details explanation of the
|
||||
implementation
|
||||
|
||||
In case of dilated_maxpool_with_argmax the returned indices from
|
||||
tf.nn.max_pool_with_argmax will point into our "no dilations" data.
|
||||
That is why they need to be mapped back to the original input data.
|
||||
It is done with `_calc_orig_argmax` function which will apply the same
|
||||
calculations, that are used in _remove_dilations when calculating the
|
||||
input data indices from output indices (check `_calc_orig_argmax` for
|
||||
detailed inline comments explaining the calculations)
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
input,
|
||||
kernel_shape,
|
||||
strides,
|
||||
dilations,
|
||||
padding="VALID",
|
||||
ceil_mode=False,
|
||||
count_include_pad=False,
|
||||
pooling_type="MAX",
|
||||
p=2):
|
||||
self.input = tf.convert_to_tensor(input)
|
||||
|
||||
self.kernel_shape = kernel_shape
|
||||
self.strides = strides
|
||||
self.dilations = dilations
|
||||
self.padding = padding
|
||||
self.is_explicit_padding = type(padding) is list
|
||||
self.ceil_mode = ceil_mode
|
||||
self.count_include_pad = count_include_pad
|
||||
self.pooling_type = pooling_type.upper()
|
||||
self.p = p
|
||||
|
||||
self.is_known_shape = self.input.shape.is_fully_defined()
|
||||
self.spatial_size = len(kernel_shape)
|
||||
self.input_rank = self.spatial_size + 2
|
||||
|
||||
# if the rank is not defined, set it to the calculated input_rank
|
||||
# rank should be known for ops like tf.gather_nd
|
||||
if not input.shape.rank:
|
||||
input.set_shape([None] * self.input_rank)
|
||||
self.orig_input_shape = tf_shape(input)
|
||||
self.input_shape = self.orig_input_shape
|
||||
|
||||
if pooling_type.startswith("MAX"):
|
||||
self.padding_constant = input.dtype.min
|
||||
else:
|
||||
self.padding_constant = 0
|
||||
|
||||
def _calc_input_ind(self, output_ind, kernel, dilation, stride):
|
||||
"""
|
||||
This function maps index from the output of _remove_dilations
|
||||
to index from the original input along single axis. It calculates
|
||||
the index inside the input data from the index of the output.
|
||||
It is used to generate the correct indexes of the values to be
|
||||
extracted by gather_nd.
|
||||
|
||||
Args:
|
||||
output_ind: vector with indices from the output to be mapped
|
||||
kernel: kernel size along the axis
|
||||
dilation: dilations along the axis
|
||||
stride: strides along the axis
|
||||
Return:
|
||||
input_ind: calculated indices
|
||||
|
||||
The formula is:
|
||||
input_ind = (output_ind // kernel) * stride +
|
||||
(output_ind % kernel) * dilation
|
||||
|
||||
Example:
|
||||
If we have following 2D input to _remove_dilations:
|
||||
[[ 0, 1, 2, 3],
|
||||
[ 4, 5, 6, 7],
|
||||
[ 8, 9, 10, 11],
|
||||
[ 12, 13, 14, 15]]
|
||||
and Kernel = [2, 2], Dilations: [2, 2], Strides: [1, 1]
|
||||
|
||||
the output of _remove_dilations will have shape [4, 4] and
|
||||
_calc_input_ind will be called twice for the two axis 0 (along
|
||||
height) and axis 1 (along width) with
|
||||
|
||||
output_ind = [0, 1, 2, 3]
|
||||
|
||||
which will result in:
|
||||
|
||||
input_ind = [0, 2, 1, 3]
|
||||
"""
|
||||
return (output_ind // kernel) * (stride - kernel * dilation) + \
|
||||
output_ind * dilation
|
||||
|
||||
def _calc_orig_argmax(self, ind):
|
||||
"""
|
||||
Map result argxmax to the original input indices
|
||||
|
||||
Maps indices generated by maxpool_with_argmax on top of the
|
||||
dilation reduced input to the orignal input indices
|
||||
"""
|
||||
|
||||
in_width = self.orig_input_shape[2]
|
||||
num_channels = self.orig_input_shape[3]
|
||||
output_width = self.output_shape[2]
|
||||
|
||||
# mod_floor op is not implemented on GPU
|
||||
# implement it using: a % b = a - (a // b) * b
|
||||
|
||||
# inRow = (ind // num_channels) // output_width
|
||||
# inCol = (ind // num_channels) % output_width
|
||||
# ind_channel = ind % num_channels
|
||||
|
||||
ind_nhw = ind // num_channels
|
||||
|
||||
inRow = ind_nhw // output_width
|
||||
inCol = ind_nhw - (ind_nhw // output_width) * output_width
|
||||
|
||||
ind_channel = ind - ind_nhw * num_channels
|
||||
|
||||
row = self._calc_input_ind(inRow, self.kernel_shape[0], self.dilations[0],
|
||||
self.strides[0]) - self.pads[0]
|
||||
col = self._calc_input_ind(inCol, self.kernel_shape[1], self.dilations[1],
|
||||
self.strides[1]) - self.pads[2]
|
||||
|
||||
new_ind = num_channels * (row * in_width + col) + ind_channel
|
||||
return new_ind
|
||||
|
||||
def _remove_dilations(self):
|
||||
"""
|
||||
This method removes the dilations by extracting the values from
|
||||
the input for every sliding window according to the dilations,
|
||||
strides and kernel size and generates output that can be used by
|
||||
pooling operations with strides = kernel_shape to accomplish
|
||||
dilated pooling
|
||||
|
||||
Example:
|
||||
Input: [[ 0, 1, 2, 3],
|
||||
[ 4, 5, 6, 7],
|
||||
[ 8, 9, 10, 11],
|
||||
[ 12, 13, 14, 15]]
|
||||
|
||||
Kernel: [2, 2]
|
||||
Dilations: [2, 2]
|
||||
Strides: [1, 1]
|
||||
|
||||
Will return:
|
||||
[[ 0, 2, 1, 3],
|
||||
[ 8, 10, 9, 11],
|
||||
[ 4, 6, 5, 7],
|
||||
[ 12, 14, 13, 15]]
|
||||
|
||||
After max_pool2d with kernel_shape = strides = [2, 2]
|
||||
the result is:
|
||||
[[ 10, 11],
|
||||
[ 14, 15]]
|
||||
"""
|
||||
|
||||
input_shape = tf_shape(self.input)
|
||||
in_spatial_shape = input_shape[1:self.spatial_size + 1]
|
||||
|
||||
channels_count = input_shape[self.spatial_size + 1]
|
||||
# Initialize gather_ind with the range of channels
|
||||
# e.g. [0 1]
|
||||
gather_ind = tf.range(channels_count, dtype=tf.int64)
|
||||
# convert the vector to column vector
|
||||
# in the following logic we use column vectors
|
||||
gather_ind = tf.expand_dims(gather_ind, 1)
|
||||
|
||||
# initilize the output_shape with zeros
|
||||
# self.output_shape will contain the shape of the
|
||||
# output tensor after the loop below is executed
|
||||
self.output_shape = [0] * (self.spatial_size + 2)
|
||||
self.output_shape[0] = input_shape[0]
|
||||
"""
|
||||
Loop over the input spatial dimensions starting from the
|
||||
last (most internal) going up to the first dimension
|
||||
|
||||
On every step of the loop calculate the output indices and
|
||||
map them to the input indices using `_calc_input_ind`,
|
||||
then "combine" with the already calculated indices from the
|
||||
previous dimensions using cartesian product.
|
||||
|
||||
For the following example input:
|
||||
|
||||
Input: [[ 0, 1, 2, 3],
|
||||
[ 4, 5, 6, 7],
|
||||
[ 8, 9, 10, 11],
|
||||
[ 12, 13, 14, 15]]
|
||||
|
||||
Kernel: [2, 2]
|
||||
Dilations: [2, 2]
|
||||
Strides: [1, 1]
|
||||
|
||||
these are the steps that will be executed:
|
||||
|
||||
1. Initilize gather_ind = [[0]] # we have only 1 channel
|
||||
|
||||
2. Loop step 0 (axis 1):
|
||||
filter_size = 3
|
||||
output_size = 4
|
||||
dim_ind = [[0]
|
||||
[2]
|
||||
[1]
|
||||
[3]]
|
||||
|
||||
gather_ind = [[0 0]
|
||||
[2 0]
|
||||
[1 0]
|
||||
[3 0]]
|
||||
|
||||
3. Loop step 1 (axis 0):
|
||||
filter_size = 3
|
||||
output_size = 4
|
||||
dim_ind = [[0]
|
||||
[2]
|
||||
[1]
|
||||
[3]]
|
||||
|
||||
gather_ind = [[0 0 0]
|
||||
[0 2 0]
|
||||
[0 1 0]
|
||||
[0 3 0]
|
||||
[2 0 0]
|
||||
[2 2 0]
|
||||
[2 1 0]
|
||||
[2 3 0]
|
||||
[1 0 0]
|
||||
[1 2 0]
|
||||
[1 1 0]
|
||||
[1 3 0]
|
||||
[3 0 0]
|
||||
[3 2 0]
|
||||
[3 1 0]
|
||||
[3 3 0]]
|
||||
|
||||
These are the indices used for gather_nd operation to collect
|
||||
the values from the input data.
|
||||
"""
|
||||
|
||||
for dim in range(self.spatial_size - 1, -1, -1):
|
||||
filter_size = (self.kernel_shape[dim] - 1) * \
|
||||
self.dilations[dim] + 1
|
||||
output_size = ((
|
||||
(in_spatial_shape[dim] - filter_size) // self.strides[dim]) + 1
|
||||
) * self.kernel_shape[dim]
|
||||
self.output_shape[dim + 1] = output_size
|
||||
|
||||
# initialize the output dimension index with the range of the
|
||||
# dimension output size (e.g. 4): [0, 1, 2, 3]
|
||||
dim_ind = tf.range(output_size)
|
||||
|
||||
# calculate the matching indices in the input data
|
||||
# [0, 1, 2, 3] will calculate to [0, 2, 1, 3]
|
||||
# from the above example
|
||||
dim_ind = self._calc_input_ind(dim_ind, self.kernel_shape[dim],
|
||||
self.dilations[dim], self.strides[dim])
|
||||
# convert to column vector
|
||||
dim_ind = tf.expand_dims(dim_ind, 1)
|
||||
|
||||
# "combine" current dimension indices with the previous dimensions
|
||||
# using cartesian product
|
||||
gather_ind = tf_product(dim_ind, gather_ind)
|
||||
|
||||
# The result from the above loop for 2D data will be:
|
||||
# [[y1, x1, c], [y2, x2, c], ..., [yn, xm, c]] where n is the height,
|
||||
# m is the width and c is the channel number.
|
||||
|
||||
# set the channels count in the output_shape
|
||||
self.output_shape[self.spatial_size + 1] = channels_count
|
||||
|
||||
# expand the dimensions to match the input dimensions + 1
|
||||
for x in range(self.spatial_size):
|
||||
gather_ind = tf.expand_dims(gather_ind, 0)
|
||||
# dublicate the indices for every batch
|
||||
gather_ind = tf.tile(gather_ind,
|
||||
[input_shape[0]] + [1] * (self.spatial_size + 1))
|
||||
|
||||
# extract the selected values from the input
|
||||
output = tf.gather_nd(self.input, gather_ind, batch_dims=1)
|
||||
# reshape the output to the correct shape calculated earlier
|
||||
output = tf.reshape(output, self.output_shape)
|
||||
|
||||
return output
|
||||
|
||||
def _calc_pads_same(self, in_spatial_shape):
|
||||
"""
|
||||
Calculate SAME_* paddings.
|
||||
"""
|
||||
|
||||
pad_ops = pooling_helper.pad_numpy_ops if self.is_known_shape else \
|
||||
pooling_helper.pad_tf_ops
|
||||
|
||||
return pooling_helper.calc_pads_same(in_spatial_shape, self.kernel_shape,
|
||||
self.strides, self.dilations,
|
||||
self.padding, pad_ops, 2)
|
||||
|
||||
def _calc_pads_explicit(self):
|
||||
"""
|
||||
Calculate explicit padding
|
||||
"""
|
||||
assert type(self.padding) is list
|
||||
|
||||
pads = []
|
||||
for i in range(self.spatial_size):
|
||||
pads += [self.padding[i], self.padding[i + self.spatial_size]]
|
||||
return pads
|
||||
|
||||
def _calc_pads_ceil_mode(self, in_spatial_shape):
|
||||
"""
|
||||
Calculate padding in ceil_mode
|
||||
"""
|
||||
|
||||
pads = []
|
||||
for i in range(self.spatial_size):
|
||||
dim_size = in_spatial_shape[i]
|
||||
filter_size = (self.kernel_shape[i] - 1) * self.dilations[i] + 1
|
||||
out_size = (dim_size - filter_size) / self.strides[i]
|
||||
if self.is_known_shape:
|
||||
pad_size = (np.ceil(out_size) - np.floor(out_size)).astype(np.int64)
|
||||
else:
|
||||
pad_size = tf.cast(
|
||||
tf.math.ceil(out_size) - tf.math.floor(out_size), tf.int64)
|
||||
|
||||
pads += [0, pad_size * self.strides[i]]
|
||||
return pads
|
||||
|
||||
def _calc_pads(self, in_spatial_shape):
|
||||
if self.is_known_shape:
|
||||
pads = np.zeros([self.spatial_size * 2], np.int64)
|
||||
else:
|
||||
pads = tf.zeros([self.spatial_size * 2], tf.int64)
|
||||
|
||||
# check for explicit padding
|
||||
if type(self.padding) is list:
|
||||
pads += self._calc_pads_explicit()
|
||||
elif self.padding.lower().startswith("same"):
|
||||
pads += self._calc_pads_same(in_spatial_shape)
|
||||
|
||||
# when padding is set to SAME, ceil_mode will not do anything
|
||||
# because output sizes will be multiple of the strides
|
||||
if self.ceil_mode and (type(self.padding) is list or
|
||||
not self.padding.lower().startswith("same")):
|
||||
new_spatial_shape = [
|
||||
in_spatial_shape[i] + pads[i * 2] + pads[i * 2 + 1]
|
||||
for i in range(self.spatial_size)
|
||||
]
|
||||
pads += self._calc_pads_ceil_mode(new_spatial_shape)
|
||||
return pads
|
||||
|
||||
def _pad_input(self):
|
||||
"""
|
||||
Pad the input according to the parameters
|
||||
"""
|
||||
# check if we need to do any padding at all
|
||||
if not self.ceil_mode and ((type(self.padding) is list and
|
||||
self.padding == [0] * self.spatial_size * 2) or
|
||||
self.padding == "VALID"):
|
||||
self.pads = np.array([0] * self.spatial_size * 2)
|
||||
return (self.input, self.pads)
|
||||
|
||||
in_spatial_shape = self.input_shape[1:self.spatial_size + 1]
|
||||
pads = self._calc_pads(in_spatial_shape)
|
||||
|
||||
if self.is_known_shape and np.count_nonzero(pads) == 0:
|
||||
self.pads = pads
|
||||
return (self.input, pads)
|
||||
|
||||
tf_paddings = [[0, 0]]
|
||||
for i in range(self.spatial_size):
|
||||
tf_paddings += [[pads[i * 2], pads[i * 2 + 1]]]
|
||||
tf_paddings += [[0, 0]]
|
||||
|
||||
self.input = tf.pad(
|
||||
self.input,
|
||||
tf_paddings,
|
||||
mode='CONSTANT',
|
||||
constant_values=self.padding_constant)
|
||||
# update input shape and pads values
|
||||
self.input_shape = tf_shape(self.input)
|
||||
self.pads = pads
|
||||
|
||||
def _calc_argmax_without_padding(self, ind):
|
||||
"""
|
||||
Calculate the original indices as they would be without padding
|
||||
"""
|
||||
in_width = self.orig_input_shape[2]
|
||||
padded_width = self.input_shape[2]
|
||||
num_channels = self.input_shape[3]
|
||||
|
||||
# mod_floor op is not implemented on GPU
|
||||
# implement it using: a % b = a - (a // b) * b
|
||||
|
||||
# ind_nhw = ind // num_channels
|
||||
# ind_channel = ind % num_channels
|
||||
|
||||
ind_nhw = ind // num_channels
|
||||
ind_channel = ind - ind_nhw * num_channels
|
||||
|
||||
new_ind = (ind_nhw // padded_width) * (self.pads[2] + self.pads[3])
|
||||
new_ind = ind_nhw - new_ind - self.pads[0] * in_width - self.pads[2]
|
||||
new_ind = num_channels * new_ind + ind_channel
|
||||
return new_ind
|
||||
|
||||
def dilated_maxpool_with_argmax(self, force_custom_impl=False):
|
||||
"""
|
||||
Do a dilated maxpool and return indices/argmax
|
||||
"""
|
||||
# Tensorflow does not support maxpool_with_argmax on
|
||||
# spatial_size != 2
|
||||
assert self.spatial_size == 2
|
||||
|
||||
if list(self.dilations) != [1] * self.spatial_size or \
|
||||
force_custom_impl:
|
||||
# pad the input
|
||||
self._pad_input()
|
||||
|
||||
new_input = self._remove_dilations()
|
||||
kernel_shape = [1] + list(self.kernel_shape) + [1]
|
||||
pooled, new_ind = tf.nn.max_pool_with_argmax(
|
||||
new_input, ksize=kernel_shape, strides=kernel_shape, padding="VALID")
|
||||
new_ind = self._calc_orig_argmax(new_ind)
|
||||
else:
|
||||
self.pads = np.array([0] * self.spatial_size * 2)
|
||||
if type(self.padding) is list or \
|
||||
self.padding.lower() == "same_lower":
|
||||
# pad the input
|
||||
self._pad_input()
|
||||
|
||||
padding_ = "VALID"
|
||||
elif self.padding.lower() == "same_upper":
|
||||
padding_ = "SAME"
|
||||
else:
|
||||
padding_ = self.padding
|
||||
|
||||
strides = [1] + list(self.strides) + [1]
|
||||
kernel_shape = [1] + list(self.kernel_shape) + [1]
|
||||
pooled, new_ind = tf.nn.max_pool_with_argmax(
|
||||
self.input, ksize=kernel_shape, strides=strides, padding=padding_)
|
||||
# if there was padding, recalculate the returned index
|
||||
# to exclude the padding
|
||||
if np.count_nonzero(self.pads) != 0:
|
||||
new_ind = self._calc_argmax_without_padding(new_ind)
|
||||
|
||||
return (pooled, new_ind)
|
||||
|
||||
def _lp_pool(self, input, ksize, strides, padding):
|
||||
window_size = np.prod(ksize)
|
||||
|
||||
input = tf.math.pow(tf.math.abs(input), self.p) * window_size
|
||||
pooled = tf.nn.avg_pool_v2(input, ksize=ksize, strides=strides,
|
||||
padding=padding)
|
||||
pooled = tf.math.pow(pooled, 1.0 / self.p)
|
||||
|
||||
return pooled
|
||||
|
||||
def dilated_pool(self, force_custom_impl=False):
|
||||
"""
|
||||
Does N-D dilated max/avg pooling. Pads the input if explicit or
|
||||
SAME_* padding is provided or ceil_mode is True
|
||||
"""
|
||||
|
||||
assert self.is_supported()
|
||||
|
||||
if self.is_explicit_padding or self.padding.lower() == "same_lower" \
|
||||
or (self.padding.lower() == "same_upper" and
|
||||
self.count_include_pad) or self.pooling_type.upper() == "LP":
|
||||
# pad the input
|
||||
self._pad_input()
|
||||
|
||||
padding_ = "VALID"
|
||||
elif self.padding.lower() == "same_upper":
|
||||
padding_ = "SAME"
|
||||
else:
|
||||
padding_ = self.padding
|
||||
|
||||
# if maxpool op with dialtions != 1 and spatial_size == 2
|
||||
# we can use tf.nn.dilation2d directly
|
||||
if self.spatial_size == 2 and self.pooling_type.startswith("MAX") \
|
||||
and self.dilations != [1] * self.spatial_size and \
|
||||
not force_custom_impl:
|
||||
strides = [1] + list(self.strides) + [1]
|
||||
dilations = [1] + list(self.dilations) + [1]
|
||||
|
||||
filter = tf.zeros(
|
||||
[self.kernel_shape[0], self.kernel_shape[1], self.input_shape[3]],
|
||||
self.input.dtype)
|
||||
pooled = tf.nn.dilation2d(
|
||||
input=self.input,
|
||||
filters=filter,
|
||||
strides=strides,
|
||||
dilations=dilations,
|
||||
padding=padding_)
|
||||
# if spatial_size < 4 and strides == 1 or dilation == 1 use tf.nn.pool
|
||||
elif self.spatial_size < 4 and (self.strides == [1] * self.spatial_size or
|
||||
self.dilations == [1] * self.spatial_size) and \
|
||||
not force_custom_impl:
|
||||
# if strides == 1 and not LpPool use tf.nn.pool directly
|
||||
if self.strides == [1] * self.spatial_size and self.pooling_type != "LP":
|
||||
pooled = tf.nn.pool(
|
||||
self.input,
|
||||
window_shape=self.kernel_shape,
|
||||
dilations=self.dilations,
|
||||
strides=self.strides,
|
||||
padding=padding_,
|
||||
pooling_type=self.pooling_type)
|
||||
else:
|
||||
# othwerwise check the pooling_type and use the correct op
|
||||
if self.pooling_type.startswith("MAX"):
|
||||
op = tf.nn.max_pool_v2
|
||||
elif self.pooling_type == "AVG":
|
||||
op = tf.nn.avg_pool_v2
|
||||
elif self.pooling_type == "LP":
|
||||
op = self._lp_pool
|
||||
else:
|
||||
raise ValueError("%d-D %s pooling is not supported." %
|
||||
(self.spatial_size, self.pooling_type))
|
||||
pooled = op(self.input, ksize=self.kernel_shape, strides=self.strides,
|
||||
padding=padding_)
|
||||
# in any other case we use custom implementation _remove_dilations
|
||||
# to reduce atrous/dilated pooling into regular pooling and selecting
|
||||
# only the values of the input that should have been selected by
|
||||
# applying the strides and dilations. Then use tf.nn.pool with
|
||||
# strides = kernel_shape and no dilations
|
||||
else:
|
||||
if padding_ == "SAME":
|
||||
# pad the input
|
||||
self._pad_input()
|
||||
input_ = self._remove_dilations()
|
||||
if self.pooling_type=="LP":
|
||||
pooled = self._lp_pool(
|
||||
input_,
|
||||
ksize=self.kernel_shape,
|
||||
strides=self.kernel_shape,
|
||||
padding="VALID")
|
||||
|
||||
else:
|
||||
pooled = tf.nn.pool(
|
||||
input_,
|
||||
window_shape=self.kernel_shape,
|
||||
strides=self.kernel_shape,
|
||||
padding="VALID",
|
||||
pooling_type=self.pooling_type)
|
||||
return pooled
|
||||
|
||||
def is_supported(self):
|
||||
"""
|
||||
Function to check if the current set of arguments are
|
||||
supported for average pool
|
||||
"""
|
||||
# check for maxpool
|
||||
if self.pooling_type.startswith("MAX") or \
|
||||
self.pooling_type=="LP":
|
||||
return True
|
||||
else:
|
||||
# if count_include_pad is true it is fully supported
|
||||
if self.count_include_pad:
|
||||
return True
|
||||
# ceil mode is not supported
|
||||
elif self.ceil_mode:
|
||||
return False
|
||||
# explicit padding with padding values set to 0 is supported
|
||||
elif (self.is_explicit_padding and
|
||||
self.padding == [0] * self.spatial_size * 2):
|
||||
return True
|
||||
# "valid" and "same_upper" auto padding is supported
|
||||
elif (not self.is_explicit_padding and
|
||||
self.padding.lower() in ["valid", "same_upper"]):
|
||||
return True
|
||||
# any other case is not supported
|
||||
else:
|
||||
return False
|
||||
Reference in New Issue
Block a user