mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-13 18:08:13 +00:00
Optimizing Upsample op (#352)
This commit is contained in:
parent
22337bb641
commit
1653ba9fcc
2 changed files with 28 additions and 21 deletions
|
|
@ -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_) {
|
||||
|
|
|
|||
|
|
@ -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_);
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue