Added Padding3D and Pooling3DAttributes.
PiperOrigin-RevId: 288367645 Change-Id: I3c37c986cfff7d22f1be50498f0ac80f878bf6f7
This commit is contained in:
parent
a376886b16
commit
77994fbf12
@ -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);
|
||||
|
@ -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,
|
||||
|
@ -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
|
||||
|
Loading…
Reference in New Issue
Block a user