Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[quant] ReflectionPad2d #48036

Closed
wants to merge 5 commits into from
Closed

[quant] ReflectionPad2d #48036

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
7e9ae57
[quant] ReflectionPad2d
z-a-f Nov 16, 2020
bc87146
Update on "[quant] ReflectionPad2d"
z-a-f Nov 16, 2020
2731a91
Update on "[quant] ReflectionPad2d"
z-a-f Nov 16, 2020
e1af264
Update on "[quant] ReflectionPad2d"
z-a-f Nov 16, 2020
7cc6b2e
Update on "[quant] ReflectionPad2d"
z-a-f Nov 16, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
61 changes: 46 additions & 15 deletions aten/src/ATen/native/ReflectionPad.cpp
View file
Open in desktop
Expand Up @@ -346,23 +346,43 @@ void reflection_pad2d_out_template(
if (input.ndimension() == 3) {
/* resize output */
output.resize_({nplane, output_h, output_w});
AT_DISPATCH_FLOATING_TYPES(input.scalar_type(), "reflection_pad2d", [&] {
reflection_pad2d_out_frame(
input.data_ptr<scalar_t>(), output.data_ptr<scalar_t>(),
nplane,
input_w, input_h, output_w, output_h,
pad_l, pad_t);
});
if (input.is_quantized()) {
AT_DISPATCH_QINT_TYPES(input.scalar_type(), "qreflection_pad2d", [&] {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

should we just add a dispatch that works for both floating types and qint types?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done!

reflection_pad2d_out_frame(
input.data_ptr<scalar_t>(), output.data_ptr<scalar_t>(),
nplane,
input_w, input_h, output_w, output_h,
pad_l, pad_t);
});
} else {
AT_DISPATCH_FLOATING_TYPES(input.scalar_type(), "reflection_pad2d", [&] {
reflection_pad2d_out_frame(
input.data_ptr<scalar_t>(), output.data_ptr<scalar_t>(),
nplane,
input_w, input_h, output_w, output_h,
pad_l, pad_t);
});
}
} else {
/* resize output */
output.resize_({nbatch, nplane, output_h, output_w});
AT_DISPATCH_FLOATING_TYPES(input.scalar_type(), "reflection_pad2d", [&] {
reflection_pad2d_out_loop(
input.data_ptr<scalar_t>(), output.data_ptr<scalar_t>(),
nbatch, nplane,
input_w, input_h, output_w, output_h,
pad_l, pad_t);
});
if (input.is_quantized()) {
AT_DISPATCH_QINT_TYPES(input.scalar_type(), "qreflection_pad2d", [&] {
reflection_pad2d_out_loop(
input.data_ptr<scalar_t>(), output.data_ptr<scalar_t>(),
nbatch, nplane,
input_w, input_h, output_w, output_h,
pad_l, pad_t);
});
} else {
AT_DISPATCH_FLOATING_TYPES(input.scalar_type(), "reflection_pad2d", [&] {
reflection_pad2d_out_loop(
input.data_ptr<scalar_t>(), output.data_ptr<scalar_t>(),
nbatch, nplane,
input_w, input_h, output_w, output_h,
pad_l, pad_t);
});
}
}
}

Expand Down Expand Up @@ -547,7 +567,18 @@ Tensor& reflection_pad2d_out_cpu(
}

Tensor reflection_pad2d_cpu(const Tensor& input, IntArrayRef padding) {
auto output = at::empty({0}, input.options());
Tensor output;
if (input.is_quantized()) {
if (input.qscheme() == kPerTensorAffine) {
output = at::_empty_affine_quantized({0}, input.options(),
input.q_scale(),
input.q_zero_point());
} else {
TORCH_CHECK(false, "Only per tensor quantization is supported");
}
} else {
output = at::empty({0}, input.options());
}
reflection_pad2d_out_template(output, input, padding);
return output;
}
Expand Down
5 changes: 2 additions & 3 deletions aten/src/ATen/native/native_functions.yaml
View file
Open in desktop
Expand Up @@ -8071,9 +8071,8 @@
use_c10_dispatcher: full
python_module: nn
dispatch:
CPU: reflection_pad1d_cpu
CPU, QuantizedCPU: reflection_pad1d_cpu
CUDA: reflection_pad1d_cuda
QuantizedCPU: reflection_pad1d_cpu

- func: reflection_pad1d_backward.grad_input(Tensor grad_output, Tensor self, int[2] padding, *, Tensor(a!) grad_input) -> Tensor(a!)
python_module: nn
Expand All @@ -8098,7 +8097,7 @@
use_c10_dispatcher: full
python_module: nn
dispatch:
CPU: reflection_pad2d_cpu
CPU, QuantizedCPU: reflection_pad2d_cpu
CUDA: reflection_pad2d_cuda

- func: reflection_pad2d_backward.grad_input(Tensor grad_output, Tensor self, int[4] padding, *, Tensor(a!) grad_input) -> Tensor(a!)
Expand Down
22 changes: 22 additions & 0 deletions test/quantization/test_quantized_op.py
View file
Open in desktop
Expand Up @@ -4111,6 +4111,28 @@ def test_reflection_pad1d(self, batch_size, channels, width, qtype):

self.assertEqual(qy_ref, qy_hat)

@given(batch_size=st.integers(1, 64),
channels=st.integers(1, 64),
height=st.integers(16, 128),
width=st.integers(16, 128),
qtype=st.sampled_from(hu._ALL_QINT_TYPES))
def test_reflection_pad2d(self, batch_size, channels, height, width, qtype):
padding = width // 4

x = torch.arange(batch_size * channels * height * width).to(torch.float)
x = x.resize(batch_size, channels, height, width)
# Per-Tensor test
scale, zp = _calculate_dynamic_qparams(x, qtype)
qx = torch.quantize_per_tensor(x, scale, zp, qtype)

padding_op = torch.nn.ReflectionPad2d(padding)

y_ref = padding_op(x)
qy_ref = torch.quantize_per_tensor(y_ref, scale, zp, qtype)
qy_hat = padding_op(qx)

self.assertEqual(qy_ref, qy_hat)

@given(batch_size=st.integers(1, 64),
channels=st.integers(1, 64),
hwd=st.integers(1, 16), # For 3D, max input size would be 16x16x16
Expand Down