Add vmap support for `torch.nn.functional.smooth_l1_loss` #98357

yhl48 · 2023-04-04T22:27:59Z

Partially fixes #97246 and #97558.

pytorch-bot · 2023-04-04T22:28:02Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/98357

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 03570dd:
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

aten/src/ATen/functorch/BatchRulesLoss.cpp

kshitij12345 · 2023-04-05T07:32:51Z

aten/src/ATen/functorch/BatchRulesLoss.cpp

@@ -282,6 +303,7 @@ TORCH_LIBRARY_IMPL(aten, FuncTorchBatched, m) {
  m.impl("nll_loss_backward", nll_loss_backward_decomposition);
  m.impl("nll_loss2d_backward", nll_loss_backward_decomposition);
  VMAP_SUPPORT(mse_loss, mse_loss_batch_rule);
+  VMAP_SUPPORT(smooth_l1_loss, smooth_l1_loss_batch_rule);


minor: Move this below the comment // mse_loss_backwards uses a decomposition for its batch rule (as it corresponds to mse_loss).

Thanks, I will reorder that in an updated commit.

kshitij12345 · 2023-04-05T08:04:34Z

runtimeerror: aten::smooth_l1_loss_backward hit the vmap fallback which is currently disabled

We need to add batching rule for smooth_l1_loss_backward.

We can proceed in either of the way

You can either add it in this PR.
Keep the relevant failing test as xfail and pursue it in a separate PR.

yhl48 · 2023-04-06T08:03:05Z

runtimeerror: aten::smooth_l1_loss_backward hit the vmap fallback which is currently disabled

We need to add batching rule for smooth_l1_loss_backward.

We can proceed in either of the way

You can either add it in this PR.

Keep the relevant failing test as xfail and pursue it in a separate PR.

I will add it in this PR, but probably need some pointer here. What I think should be done are:

Implement a function smooth_l1_loss_backward
Add that function to TORCH_LIBRARY_IMPL

Is there an easier way to do it? Perhaps something similar to mse_loss_backwards as commented in the code // mse_loss_backwards uses a decomposition for its batch rule ?

@kshitij12345

kshitij12345 · 2023-04-10T05:45:23Z

Apologies for the delayed response.

For adding a decomposition, you'll need to add it in the following file.

pytorch/torch/_decomp/decompositions.py

Lines 356 to 362 in 95621b3

    
           @register_decomposition(aten.mse_loss_backward) 
        
           @pw_cast_for_opmath 
        
           def mse_loss_backward( 
        
               grad_output: Tensor, input: Tensor, target: Tensor, reduction: int 
        
           ): 
        
               norm = 2.0 / input.numel() if reduction == Reduction.MEAN.value else 2.0 
        
               return norm * (input - target) * grad_output

And then register this decomposition with vmap, similar to this PR https://github.com/pytorch/functorch/pull/866/files

However, I would recommend doing that in a separate PR.

yhl48 · 2023-04-11T08:47:51Z

@kshitij12345 Thanks! I haven't had the chance to wrap this PR up, but I will do it the soonest I can, and send a separate PR for smooth_l1_loss_backward

This reverts commit b7817ea.

yhl48 · 2023-04-13T14:12:59Z

@kshitij12345 done!

aten/src/ATen/functorch/BatchRulesLoss.cpp

kshitij12345

Overall looks good! Just minor changes related to naming of lambda arguments and cleaning up commented code :)

Thanks!

aten/src/ATen/functorch/BatchRulesLoss.cpp

yhl48 · 2023-04-14T08:39:48Z

Thanks, appreciate your help throughout!

kshitij12345

LGTM! Thank you @yhl48

kshitij12345 · 2023-04-14T09:15:07Z

@pytorchbot merge

pytorchmergebot · 2023-04-14T09:17:13Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

Partially fixes #97246 and #97558. Pull Request resolved: #98357 Approved by: https://github.com/kshitij12345

Follow-up of #98357 Pull Request resolved: #99429 Approved by: https://github.com/kshitij12345, https://github.com/zou3519

yhl48 requested review from zou3519, Chillee and kshitij12345 as code owners April 4, 2023 22:27

yhl48 changed the title ~~Batch smooth l1 loss~~ Add vmap support for smooth_l1_loss Apr 4, 2023

pytorchbot added the open source label Apr 4, 2023

yhl48 force-pushed the batch_smooth_l1_loss branch from 95d48c2 to 6d2f9c4 Compare April 4, 2023 23:06

yhl48 requested review from mruberry and ngimel as code owners April 4, 2023 23:06

yhl48 marked this pull request as draft April 4, 2023 23:07

yhl48 force-pushed the batch_smooth_l1_loss branch from 6d2f9c4 to 7a13a9b Compare April 4, 2023 23:16

yhl48 added 2 commits April 4, 2023 23:21

Added vmap support for torch.nn.functional.smooth_l1_loss

0c43d23

Removed xfail

b7817ea

yhl48 force-pushed the batch_smooth_l1_loss branch from ecab6ac to b7817ea Compare April 4, 2023 23:21

yhl48 changed the title ~~Add vmap support for smooth_l1_loss~~ Add vmap support for torch.nn.functional.smooth_l1_loss Apr 4, 2023

kshitij12345 reviewed Apr 5, 2023

View reviewed changes

aten/src/ATen/functorch/BatchRulesLoss.cpp Show resolved Hide resolved

kshitij12345 reviewed Apr 5, 2023

View reviewed changes

yhl48 added 3 commits April 13, 2023 09:43

Added a helper function loss_batch_rule_helper

3b1865f

Revert "Removed xfail"

20bea84

This reverts commit b7817ea.

Removed xfail for nn.functional.smooth_l1_loss

a8d3e09

yhl48 marked this pull request as ready for review April 13, 2023 14:12

kshitij12345 reviewed Apr 13, 2023

View reviewed changes

aten/src/ATen/functorch/BatchRulesLoss.cpp Outdated Show resolved Hide resolved

dagitses added the triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module label Apr 13, 2023

Use template instead of std::function

b6a93fd

kshitij12345 requested changes Apr 14, 2023

View reviewed changes

aten/src/ATen/functorch/BatchRulesLoss.cpp Outdated Show resolved Hide resolved

aten/src/ATen/functorch/BatchRulesLoss.cpp Outdated Show resolved Hide resolved

Cleaned up code and renamed arguments

03570dd

kshitij12345 approved these changes Apr 14, 2023

View reviewed changes

kshitij12345 added the release notes: functorch release notes category; Pertaining to torch.func or pytorch/functorch label Apr 14, 2023

pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Apr 14, 2023

pytorchmergebot added the merging label Apr 14, 2023

pytorchmergebot added the Merged label Apr 14, 2023

pytorchmergebot closed this in 298cc5c Apr 14, 2023

yhl48 mentioned this pull request Apr 18, 2023

Add vmap support for smooth_l1_loss_backward #99429

Closed

ZainRizvi pushed a commit that referenced this pull request Apr 19, 2023

Add vmap support for torch.nn.functional.smooth_l1_loss (#98357)

dc7cdc6

Partially fixes #97246 and #97558. Pull Request resolved: #98357 Approved by: https://github.com/kshitij12345

pytorchmergebot pushed a commit that referenced this pull request Apr 28, 2023

Add vmap support for smooth_l1_loss_backward (#99429)

07c02b9

Follow-up of #98357 Pull Request resolved: #99429 Approved by: https://github.com/kshitij12345, https://github.com/zou3519

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add vmap support for `torch.nn.functional.smooth_l1_loss` #98357

Add vmap support for `torch.nn.functional.smooth_l1_loss` #98357

yhl48 commented Apr 4, 2023

pytorch-bot bot commented Apr 4, 2023 •

edited

kshitij12345 Apr 5, 2023

yhl48 Apr 5, 2023

kshitij12345 commented Apr 5, 2023

yhl48 commented Apr 6, 2023 •

edited

kshitij12345 commented Apr 10, 2023

yhl48 commented Apr 11, 2023

yhl48 commented Apr 13, 2023

kshitij12345 left a comment

yhl48 commented Apr 14, 2023

kshitij12345 left a comment

kshitij12345 commented Apr 14, 2023

pytorchmergebot commented Apr 14, 2023

Add vmap support for torch.nn.functional.smooth_l1_loss #98357

Add vmap support for torch.nn.functional.smooth_l1_loss #98357

Conversation

yhl48 commented Apr 4, 2023

pytorch-bot bot commented Apr 4, 2023 • edited

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/98357

✅ No Failures

kshitij12345 Apr 5, 2023

Choose a reason for hiding this comment

yhl48 Apr 5, 2023

Choose a reason for hiding this comment

kshitij12345 commented Apr 5, 2023

yhl48 commented Apr 6, 2023 • edited

kshitij12345 commented Apr 10, 2023

yhl48 commented Apr 11, 2023

yhl48 commented Apr 13, 2023

kshitij12345 left a comment

Choose a reason for hiding this comment

yhl48 commented Apr 14, 2023

kshitij12345 left a comment

Choose a reason for hiding this comment

kshitij12345 commented Apr 14, 2023

pytorchmergebot commented Apr 14, 2023

Merge started

Add vmap support for `torch.nn.functional.smooth_l1_loss` #98357

Add vmap support for `torch.nn.functional.smooth_l1_loss` #98357

pytorch-bot bot commented Apr 4, 2023 •

edited

yhl48 commented Apr 6, 2023 •

edited