Optimizing Upsample op (#352)

This commit is contained in:
Du Li 2019-01-18 16:36:00 -08:00 committed by GitHub
parent 22337bb641
commit 1653ba9fcc
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 28 additions and 21 deletions

View file

@ -22,13 +22,14 @@ ONNX_CPU_OPERATOR_VERSIONED_TYPED_KERNEL(
KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType<int32_t>()),
Upsample<int32_t>);
void upsampleNearest2x(
template <typename T>
void UpsampleNearest2x(
int64_t batch_size,
int64_t num_channels,
int64_t input_height,
int64_t input_width,
const float* input,
float* output) {
const T* input,
T* output) {
const int64_t output_height = input_height * 2;
const int64_t output_width = input_width * 2;
for (int64_t n = 0; n < batch_size; ++n) {
@ -36,7 +37,7 @@ void upsampleNearest2x(
for (int64_t y = 0; y < output_height; ++y) {
const int64_t in_y = y / 2;
for (int64_t x = 0; x < input_width; ++x) {
const float v = input[in_y * input_width + x];
const T v = input[in_y * input_width + x];
const int64_t oidx = output_width * y + x * 2;
output[oidx + 0] = v;
output[oidx + 1] = v;
@ -49,7 +50,7 @@ void upsampleNearest2x(
}
template <typename T>
Status upsampleNearest(const T* input,
Status UpsampleNearest(const T* input,
T* output,
const TensorShape& input_shape,
const TensorShape& output_shape,
@ -59,19 +60,23 @@ Status upsampleNearest(const T* input,
if (input_shape.NumDimensions() != output_shape.NumDimensions())
return Status(ONNXRUNTIME, FAIL, "Upsample: input/output value's dimension mismatch");
auto n_dim = input_shape.NumDimensions();
for (size_t i = 0, size = output_shape.Size(); i < size; i++) {
size_t old_idx = 0;
size_t cur_idx = i;
if (scales.size() == 4 && scales[0] == 1 && scales[1] == 1 && scales[2] == 2 && scales[3] == 2) {
UpsampleNearest2x<T>(input_shape[0], input_shape[1], input_shape[2], input_shape[3], input, output);
} else {
for (size_t i = 0, size = output_shape.Size(); i < size; i++) {
size_t old_idx = 0;
size_t cur_idx = i;
int64_t base = 1;
for (int64_t j = static_cast<int64_t>(n_dim - 1); j >= 0; j--) {
auto tmp = cur_idx % output_shape[j];
old_idx += (std::min(static_cast<int64_t>(tmp / scales[j]), input_shape[j] - 1)) * base;
base *= input_shape[j];
cur_idx /= output_shape[j];
int64_t base = 1;
for (int64_t j = static_cast<int64_t>(n_dim - 1); j >= 0; j--) {
auto tmp = cur_idx % output_shape[j];
old_idx += (std::min(static_cast<int64_t>(tmp / scales[j]), input_shape[j] - 1)) * base;
base *= input_shape[j];
cur_idx /= output_shape[j];
}
output[i] = input[old_idx];
}
output[i] = input[old_idx];
}
return Status::OK();
}
@ -208,7 +213,7 @@ Status Upsample<T>::BaseCompute(OpKernelContext* context, const std::vector<floa
switch (mode_) {
case UpsampleMode::NN:
return upsampleNearest<T>(X->template Data<T>(), Y->template MutableData<T>(), X->Shape(), Y->Shape(), scales);
return UpsampleNearest<T>(X->template Data<T>(), Y->template MutableData<T>(), X->Shape(), Y->Shape(), scales);
case UpsampleMode::LINEAR: {
//What's the correct behavior of linear mode is not clear right now,
//Only support bilinear with 4D tensor to keep consistent with previous behavior
@ -227,7 +232,6 @@ Status Upsample<T>::BaseCompute(OpKernelContext* context, const std::vector<floa
}
}
template <typename T>
Status Upsample<T>::Compute(OpKernelContext* context) const {
if (OpKernel::Node().InputDefs().size() == 1 || scales_cached_) {

View file

@ -17,7 +17,7 @@ enum UpsampleMode {
class UpsampleBase {
protected:
UpsampleBase(OpKernelInfo info): scales_cached_(false) {
UpsampleBase(OpKernelInfo info) : scales_cached_(false) {
std::string mode;
ORT_ENFORCE(info.GetAttr<std::string>("mode", &mode).IsOK());
mode_ = StringToUpsampleMode(mode);
@ -51,7 +51,7 @@ class UpsampleBase {
return UpsampleMode::LINEAR;
} else {
ORT_THROW("mode attribute is " + mode + ". It can only be " +
UpsampleModeNN + "(default) or " + UpsampleModeLinear + ".");
UpsampleModeNN + "(default) or " + UpsampleModeLinear + ".");
}
}
@ -63,7 +63,7 @@ class UpsampleBase {
if (UpsampleMode::LINEAR == mode) {
ORT_ENFORCE(scales.size() == 4, "Upsample: linear mode upsample only support bilinear with 4 dimension.");
ORT_ENFORCE(((scales[0] == 1) && (scales[1] == 1)),
"Upsample: linear mode upsample only support bilinear, the first 2 scales should be 1.");
"Upsample: linear mode upsample only support bilinear, the first 2 scales should be 1.");
}
}
@ -71,6 +71,9 @@ class UpsampleBase {
const float* scale_data = scale->template Data<float>();
int64_t scales_size = scale->Shape().Size();
ORT_ENFORCE(scales_size > 0, "scales size should be greater than 0.");
if (scales.size() == 0) {
scales.resize(scales_size);
}
memcpy(scales.data(), scale_data, scales_size * sizeof(float));
ScalesValidation(scales, mode_);
}