Added Padding3D and Pooling3DAttributes.

PiperOrigin-RevId: 288367645
Change-Id: I3c37c986cfff7d22f1be50498f0ac80f878bf6f7
This commit is contained in:
A. Unique TensorFlower 2020-01-06 13:45:57 -08:00 committed by TensorFlower Gardener
parent a376886b16
commit 77994fbf12
3 changed files with 107 additions and 0 deletions

View File

@ -44,6 +44,28 @@ Padding2D& Padding2D::operator-(const Padding2D& value) {
return *this;
}
Padding3D& Padding3D::operator=(const Padding3D& value) {
prepended = value.prepended;
appended = value.appended;
return *this;
}
bool Padding3D::operator==(const Padding3D& value) {
return this->prepended == value.prepended && this->appended == value.appended;
}
bool Padding3D::operator!=(const Padding3D& value) { return !(*this == value); }
Padding3D& Padding3D::operator-(const Padding3D& value) {
prepended.h -= value.prepended.h;
prepended.w -= value.prepended.w;
prepended.d -= value.prepended.d;
appended.h -= value.appended.h;
appended.w -= value.appended.w;
appended.d -= value.appended.d;
return *this;
}
std::string ToString(enum OperationType op) {
switch (op) {
case OperationType::ABS:
@ -206,6 +228,15 @@ int32_t CalculateOutputWithoutStrides(const BHWC& input,
/*dilation=*/1);
}
template <Axis T>
int32_t CalculateOutputWithoutStrides(const BHWDC& input,
const Pooling3DAttributes& attr) {
return CalculateOutputSizeBeforeStrides(
input.get<T>(), attr.kernel.get<T>(),
attr.padding.prepended.get<T>() + attr.padding.appended.get<T>(),
/*dilation=*/1);
}
template <Axis T>
int32_t CalculateOutput(const BHWC& input,
const ConvolutionTransposedAttributes& attr) {
@ -224,6 +255,12 @@ int32_t CalculateOutput(const BHWC& input, const AttrT& attr) {
attr.strides.template get<AxisT>());
}
template <Axis AxisT, typename AttrT>
int32_t CalculateOutput(const BHWDC& input, const AttrT& attr) {
return StridedSize(CalculateOutputWithoutStrides<AxisT>(input, attr),
attr.strides.template get<AxisT>());
}
int32_t CalculateSamePadding(int32_t input, int32_t kernel, int32_t dilation,
int32_t stride) {
const int32_t dilated_kernel = (kernel - 1) * dilation + 1;
@ -255,6 +292,13 @@ int32_t CalculateSamePadding(const BHWC& input,
/*dilation=*/1, attr.strides.get<AxisT>());
}
template <Axis AxisT>
int32_t CalculateSamePadding(const BHWDC& input,
const Pooling3DAttributes& attr) {
return CalculateSamePadding(input.get<AxisT>(), attr.kernel.get<AxisT>(),
/*dilation=*/1, attr.strides.get<AxisT>());
}
template <Axis AxisT>
int32_t CalculateSamePadding(const BHWC& input,
const MaxUnpooling2DAttributes& attr) {
@ -285,6 +329,21 @@ Padding2D MakeSamePadding(const BHWC& input, const AttrT& attr) {
return padding;
}
// If padding depends on input, convert it into fixed padding.
template <class AttrT>
Padding3D MakeSamePadding(const BHWDC& input, const AttrT& attr) {
int32_t padding_height = CalculateSamePadding<Axis::HEIGHT>(input, attr);
int32_t padding_width = CalculateSamePadding<Axis::WIDTH>(input, attr);
int32_t padding_depth = CalculateSamePadding<Axis::DEPTH>(input, attr);
Padding3D padding;
padding.prepended =
HWD(padding_height / 2, padding_width / 2, padding_depth / 2);
padding.appended =
HWD(padding_height - padding_height / 2,
padding_width - padding_width / 2, padding_depth - padding_depth / 2);
return padding;
}
} // namespace
BHWC CalculateOutputShape(const BHWC& input,
@ -302,6 +361,13 @@ BHWC CalculateOutputShape(const BHWC& input, const Pooling2DAttributes& attr) {
CalculateOutput<Axis::WIDTH>(input, attr), input.c);
}
BHWDC CalculateOutputShape(const BHWDC& input,
const Pooling3DAttributes& attr) {
return BHWDC(input.b, CalculateOutput<Axis::HEIGHT>(input, attr),
CalculateOutput<Axis::WIDTH>(input, attr),
CalculateOutput<Axis::DEPTH>(input, attr), input.c);
}
BHWC CalculateOutputShape(const BHWC& input,
const Convolution2DAttributes& attr) {
return BHWC(input.b, CalculateOutput<Axis::HEIGHT>(input, attr),
@ -405,6 +471,11 @@ Padding2D CalculateSamePadding(const BHWC& input,
return MakeSamePadding(input, attr);
}
Padding3D CalculateSamePadding(const BHWDC& input,
const Pooling3DAttributes& attr) {
return MakeSamePadding(input, attr);
}
Padding2D CalculateSamePadding(const BHWC& input,
const MaxUnpooling2DAttributes& attr) {
return MakeSamePadding(input, attr);

View File

@ -92,6 +92,20 @@ struct Padding2D {
HW appended = HW(-1, -1);
};
struct Padding3D {
Padding3D() = default;
Padding3D& operator=(const Padding3D& value);
bool operator==(const Padding3D& value);
bool operator!=(const Padding3D& value);
Padding3D& operator-(const Padding3D& value);
// Padding values for every axis (if needed), where 'prepended' defines
// padding for the beginning of each axis and 'appended' represents end part
// of the corresponding axis.
HWD prepended = HWD(0, 0, 0);
HWD appended = HWD(0, 0, 0);
};
struct Crop2D : public Padding2D {};
struct SpaceToBatchAttributes {
@ -126,6 +140,18 @@ struct Pooling2DAttributes {
bool output_indices = false;
};
struct Pooling3DAttributes {
PoolingType type = PoolingType::UNDEFINED;
// Strides for every axis.
HWD strides = HWD(0, 0, 0);
HWD kernel = HWD(0, 0, 0);
Padding3D padding;
// NOTE(akulik): technically the number of outputs from Pooling node indicates
// whether indices are needed or not, but I decided to keep it inside
// attributes to simplify processing.
bool output_indices = false;
};
struct MaxUnpooling2DAttributes {
// Strides for every axis.
HW strides = HW(-1, -1);
@ -147,6 +173,10 @@ BHWC CalculateOutputShape(const BHWC& input,
// input.
BHWC CalculateOutputShape(const BHWC& input, const Pooling2DAttributes& attr);
// @return shape of a tensor after Pooling3D operation is applied to the given
// input.
BHWDC CalculateOutputShape(const BHWDC& input, const Pooling3DAttributes& attr);
// @return shape of a tensor after Concat operation is applied to the given
// input.
Status CalculateOutputShape(const std::vector<BHWC>& input,
@ -157,6 +187,11 @@ Status CalculateOutputShape(const std::vector<BHWC>& input,
Padding2D CalculateSamePadding(const BHWC& input,
const Pooling2DAttributes& attr);
// @return padding for pooling operation to make sure output keep the same shape
// as the given input.
Padding3D CalculateSamePadding(const BHWDC& input,
const Pooling3DAttributes& attr);
// @return padding for max unpooling operation to make sure output keep the same
// shape as the given input.
Padding2D CalculateSamePadding(const BHWC& input,

View File

@ -86,6 +86,7 @@ Tensor<ShapeT, Type> MakeZeroTensor(const ShapeT& shape) {
}
using TensorFloat32 = Tensor<BHWC, DataType::FLOAT32>;
using Tensor5DFloat32 = Tensor<BHWDC, DataType::FLOAT32>;
} // namespace gpu
} // namespace tflite