Refactor: Split conv backward to allow conditional gradient computation

[AsherJingkongChen]

Checklist

• Confirmed that run-checks all script has been executed.
• Made sure the book is up to date with changes in this PR.

Related Issues/PRs

Not found.

Changes

In crate burn-tensor, I split each of conv[1-3]d_backward and their transposed variants into three functions suffixed with *_x_backward, *_weight_backward, and *_bias_backward.

In crate burn-autodiff, I changed the backward implementations of conv ops to allow gradients being computed if required.

The public APIs appear to remain unchanged, since Autodiff backend (decorator) hides the changes.

Testing

All tests are passed on Apple M2 Pro. I did not add any new test because the original tests should be adequate.

[AsherJingkongChen]

Use case

This is M-SSIM loss function. I use this to compare the performance between the main and my PR. The changed version runs faster, because no grad computation is required for Conv2d.weight here.

• main (6f0e61a): 930.839125ms
• this PR (8d0852f): 41.750209ms

https://gist.github.com/AsherJingkongChen/9baa00fc1d0503e55a358c55e393fc87

[codecov]

Codecov Report

Attention: Patch coverage is 86.24754% with 70 lines in your changes missing coverage. Please review.

Project coverage is 85.92%. Comparing base (6f0e61a) to head (8d0852f).
Report is 5 commits behind head on main.

Files with missing lines	Patch %	Lines
crates/burn-autodiff/src/ops/module.rs	61.45%	69 Missing ⚠️
crates/burn-tensor/src/tensor/ops/modules/conv.rs	99.55%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2278      +/-   ##
==========================================
- Coverage   85.92%   85.92%   -0.01%     
==========================================
  Files         750      750              
  Lines       94357    94573     +216     
==========================================
+ Hits        81076    81259     +183     
- Misses      13281    13314      +33     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[nathanielsimard]

Do we need the WithBias and NoBias variant of each convolution?

[nathanielsimard]

@

Do we need the WithBias and NoBias variant of each convolution?

Because of the prepare method in the backward pass, it's not yet possible to have only one backward struct.

[AsherJingkongChen]

@

Do we need the WithBias and NoBias variant of each convolution?

Because of the prepare method in the backward pass, it's not yet possible to have only one backward struct.

@nathanielsimard I think so.