Build the affine/perspective sampling grid in at least float32 - #9658
Open
AbdulAliMamnun wants to merge 3 commits into
Open
Build the affine/perspective sampling grid in at least float32#9658AbdulAliMamnun wants to merge 3 commits into
AbdulAliMamnun wants to merge 3 commits into
Conversation
_apply_grid_transform casts the image to the grid's dtype, samples, then
casts back. It decided whether to round on the way back with a single
flag:
fp = img.dtype == grid.dtype
...
img = float_img.round_().to(img.dtype) if not fp else float_img
That flag conflates two unrelated questions: "did the image need a dtype
cast?" and "is the image an integer type that must be rounded before the
cast back?". Today the two happen to coincide, because the only images
that get cast are integer ones - a float image is always handed a grid of
its own dtype, so it takes the `fp` branch and is never rounded.
That coincidence breaks as soon as the grid is built at a wider dtype
than the image. A float16 image would then land in the `not fp` branch
and have round_() applied to pixel data in [0, 1], quantising every
pixel to exactly 0.0 or 1.0. None of the 909 existing rotate/perspective
tests catch that, because none of them exercise half precision.
Key the decision off the input dtype instead, using the explicit
integer-dtype tuple that the v1 path already uses in _cast_squeeze_out.
This is a no-op today: integer images round exactly as before and float
images are still returned untouched. It is a prerequisite for promoting
the sampling grid to float32, which is done in the following commit.
The tuple is spelled out rather than using torch.dtype.is_floating_point
because TorchScript models torch.dtype as an int, and scripting
rotate_image/affine_image/perspective_image would fail on the attribute.
Tests cover the dtype round trip for float16, float32, float64 and uint8,
that a float image is not quantised to 0.0/1.0, and that a uint8 image is
still rounded rather than truncated before the cast back.
Signed-off-by: Abdul Ali Mamnun <aamamnun@gmail.com>
Rotating a 1024x1024 float16 image produced 6716 NaN and 11 Inf pixels. After this change it produces none. Fixes pytorch#9029. The sampling grid was allocated at the incoming dtype: base_grid = torch.empty(1, oh, ow, 3, dtype=theta.dtype, ...) so for a float16 or bfloat16 image the whole grid - the bmm that applies theta, the normalisation, and the returned tensor - ran and was stored in half precision. Half precision does not have the mantissa bits to address a large image. Grid coordinates are normalised to [-1, 1], where float16 has roughly 11 bits of resolution, far fewer than the ~19 bits a 512x512 output needs to give every pixel its own source location. The visible symptom is not a large numerical error, it is a collapse in the number of *distinct* sample positions. At 512x512 rotated 30 degrees the grid carries 262144 distinct x coordinates in float32 and 16398 in float16, so blocks of neighbouring output pixels all read from an identical source location and the result comes out blocky. Perspective additionally divides by a half precision denominator, which is where the NaN and Inf above come from. Build and return the grid at torch.promote_types(<current dtype>, torch.float32). Promotion rather than a hardcoded float32 so that a float64 image keeps a float64 grid instead of being silently downcast. For float32 and float64 callers this changes nothing: promote_types is the identity there, and 1248 rotate/affine/perspective outputs spanning both v1 and v2, two sizes, nearest and bilinear, three fill variants, expand on and off, and float32/float64/uint8/int64 are bit-for-bit identical to the previous implementation. Half precision images now sample through a float32 grid, so grid_sample runs at float32 and the image is upcast for the duration of the call. Measured on CPU for rotate on a 1024x1024 image, single threaded, over 12 interleaved A/B rounds: float16 costs +20% wall time (7.0 -> 8.5 ms) and about +10 MiB peak RSS, while bfloat16 gets 21% *faster* (11.1 -> 8.8 ms) because CPU grid_sample is roughly 2x slower at bfloat16 than at float32. float32 and float64 are unchanged to within measurement noise (+-4%). Tests assert that the number of distinct sample positions in float16 is within 1% of float32 for rotate and perspective across both v1 and v2, that a float64 grid is not downcast, that the output is finite, and that the float16 result tracks the float32 result. They fail without this change. Signed-off-by: Abdul Ali Mamnun <aamamnun@gmail.com>
Promoting the grid alone does not fix bfloat16. The grid function
receives theta already built at the image's dtype, and bfloat16 carries
only 7 mantissa bits, so the transform coefficients are quantised before
the grid is ever allocated. At 512x512 rotated 30 degrees that leaves
82496 of 262144 distinct x coordinates - a 3x improvement on the 16398
of the unfixed code, but still a third of what float32 gives.
The loss is entirely caller-side, not in the grid: feeding a
bfloat16-rounded theta into the promoted float32 grid produces exactly
the same 82496 distinct coordinates. Build theta at the grid's precision
too and the count reaches the full 262144, matching float32 exactly for
both float16 and bfloat16 across rotate, affine and perspective.
The call sites computed the dtype as
dtype = image.dtype if torch.is_floating_point(image) else torch.float32
which torch.promote_types(image.dtype, torch.float32) reproduces exactly
for uint8, int8, int16, int32, int64, bool, float32 and float64. Only
float16 and bfloat16 differ, so this is a drop-in replacement at the six
grid call sites. The other uses of that expression - adjust_contrast,
gaussian_blur, _blurred_degenerate_image, autocontrast and elastic - are
deliberately left alone; they do not build a sampling grid.
This commit is self-contained and can be dropped to scope the fix to
float16 only. Reverting it restores the previous theta dtype and the
float16-only test parametrisation, leaving the grid promotion and the
cast-back rounding guard intact.
Signed-off-by: Abdul Ali Mamnun <aamamnun@gmail.com>
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/vision/9658
Note: Links to docs will display an error until the docs builds have been completed. This comment was automatically generated by Dr. CI and updates every 15 minutes. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #9029. Full analysis in
https://github.com/pytorch/vision/issues/9029#issuecomment-5554773857—
summary here.
(Disclosure: I used Claude to help with the investigation and
implementation. The analysis and design decisions are mine, and I've
verified every number below directly.)
Unpatched, a float16 1024×1024
rotateemits 6716 NaNs and 11 Infs, andresolves 262144 sample positions down to 16398 distinct coordinates. The
grid is allocated at
theta.dtype, so thebmmthat builds it runs inhalf precision.
Commits
Three, each green in isolation and each tested at its own commit.
1 — Key the cast-back rounding off the input dtype (v2
_apply_grid_transform). A no-op today, and a prerequisite for commit 2:fp = img.dtype == grid.dtypedecides both "needs a cast" and "is aninteger type needing rounding", so a promoted grid would send float16
down the rounding branch and snap every pixel to 0.0 or 1.0. All 909
existing rotate/perspective tests pass in that broken state. v1 is safe
because
_cast_squeeze_outgates rounding behind an explicitinteger-dtype tuple; this adds the same guard to v2.
2 — Build the grid at
torch.promote_types(dtype, torch.float32)atthe four grid sites across v1 and v2. Fixes float16.
3 — Promote
thetaat six call sites. bfloat16 only — at 7 mantissabits the quantized theta is itself coarse enough to collapse
coordinates, so the grid promotion alone leaves bf16 at 82496/262144.
Droppable: reverts cleanly with no conflicts and leaves the tree
byte-identical to commit 2, if you'd rather scope this to float16.
float32 / float64 are untouched
promote_types(d, float32)reproduces the existingimg.dtype if is_floating_point else float32exactly for every dtypeexcept fp16/bf16. Verified bitwise rather than by inspection: 312 outputs
per dtype across rotate/affine/perspective × v1/v2 × two sizes ×
interpolation modes × fill variants × angles × shear/translate, all
torch.equalidentical including dtype and shape.Cost
1024×1024, CPU, 12 interleaved A/B rounds:
float16 costs ~1.8 ms per megapixel — roughly half in the slower fp32
grid_sample, half in the two casts — and ~10 MiB peak RSS. bfloat16gets faster: CPU
grid_sampleat bf16 is about 2× slower than at fp32(11.37 → 5.95 ms), so the upcast more than pays for itself.
Worth noting uint8 images already take the fp32 grid path today, since
thetais float32 for any non-floating-point input. This makes halfprecision consistent with the common case rather than introducing a new
cost class.
Tests
46 new, 20 of which fail on unpatched source (verified by stashing
torchvision/and keeping the tests). They assert grid resolution,output vs float32, finiteness, dtype preservation, and that uint8 is
still rounded rather than truncated. They pass under shuffled ordering,
which matters because
TestGridPrecisionmonkeypatches_apply_grid_transformand restores it in afinally.torch.jit.scriptstill passes onrotate_image,affine_image,perspective_image,elastic_image,_apply_grid_transform,_affine_grid,_perspective_gridand the v1 equivalents — an earlierdraft used
dtype.is_floating_point, which broke 177 tests becauseTorchScript models
torch.dtypeasint.mypy --config-file mypy.inireports 7 errors in 5 files, byte-identicalto
main; none in the changed files.ufmt formatwith the pinnedufmt 1.3.3 / black 22.3.0 / usort 1.0.2 makes no changes.
flake8 torchvisionclean.test_transforms.py: 1335 passed, 0 failed.test_transforms_v2.py: 46failed, identical set before and after — all
TestJPEG, from my localbuild having no libjpeg.