[FEATURE] Add feature of attach_grad to nonleaf variables in HybridizedBlock.#20559
[FEATURE] Add feature of attach_grad to nonleaf variables in HybridizedBlock.#20559KexinFeng wants to merge 0 commit intoapache:masterfrom
Conversation
|
Hey @KexinFeng , Thanks for submitting the PR
CI supported jobs: [sanity, windows-gpu, clang, edge, unix-cpu, windows-cpu, unix-gpu, centos-cpu, website, miscellaneous, centos-gpu] Note: |
mseth10
added
the
pr-awaiting-testing
mseth10
added
pr-work-in-progress
mseth10
added
pr-awaiting-testing
python/mxnet/gluon/block.py
Outdated
| if not self._active: | ||
| raise RuntimeError('Hybridize must be active in order to use mark_vars') |
There was a problem hiding this comment.
For a good user experience, it would be great if we have an API that supports both hybrid and imperative mode. It doesn't need to use the same backend implementation, but it would be good to have a consistent frontend API.
There was a problem hiding this comment.
Now both hybridized and unhybrized are unified. This is shown in the last two examples in unittest/test_autograd.py
mseth10
added
pr-awaiting-testing
mseth10
added
pr-awaiting-testing
|
@mxnet-bot run ci [unix-gpu] |
|
Jenkins CI successfully triggered : [unix-gpu] |
mseth10
added
pr-awaiting-testing
|
@mxnet-bot run ci [website] |
|
Jenkins CI successfully triggered : [website] |
mseth10
added
pr-awaiting-testing
|
@mxnet-bot run ci [website] |
|
Jenkins CI successfully triggered : [website] |
|
@mxnet-bot run ci [linkcheck] |
|
None of the jobs entered are supported. |
|
@mxnet-bot run ci [centos-gpu] |
|
Jenkins CI successfully triggered : [centos-gpu] |
src/imperative/imperative_utils.cc
Outdated
| int mark_id = std::stoi(it->second); | ||
| CHECK_LT(mark_id, nleafs.size()) | ||
| << "Mark_id exceeds the nonleaf list size."; | ||
| nleafs[mark_id]->copy_autograd_entry_(ndoutputs[0]); |
There was a problem hiding this comment.
Have you verified this is correct if a marked variable has multiple outputs? I see your test case below only considers variables with a single output:
def forward(self, a, b):
out1 = a*b
out2 = out1 * a
self.mark_vars(out1)
return out2
There was a problem hiding this comment.
I have added the unit test which considers multiple outputs. This is in the commit 'multiple_outputs', also shown below. It shows that it works as expected.
def forward(self, a, b, c):
out1 = self.intermediate(('out1_0', 'out1_1'), ((a+b)*c, a*b), grad_req='write')
out2 = self.intermediate('out2', out1[1] * a)
return out2
In the above commented part, in line 177, indeed the use of ndoutputs[0] assumes that in the forward graph each computation node has only one output. This is consistent with line 152. Based on the test above, for now it seems that this assumption works for the multiple output case: ((a+b)*c, a*b).
|
@mxnet-bot run ci [centos-gpu] |
|
Jenkins CI successfully triggered : [centos-gpu] |
5 participants
Description
The PR adds the support for fetching the gradients of intermediate variables in a gluon HybridizedBlock. This applies uniformly to both when
block.hybridize()is on and off. This generates theattach_gradimplemented in implemented in PR#20500.The motivation of this feature comes from this issue#11865.
Checklist
Essentials
Changes
block.pywheremark_varsandget_mark_varsare added along withMXNDArrayMarkDCVariables.cached_op.invokein cpp backend andCachedOp.__call__have been editted to include the pass of marked nonleaf ndarrays.set_nleafsmethod is added intoCachedOpclass to store the marked nonleaf ndarrays.void RunGraph, marked nonleaf ndarrays are linked to the marked computational node for autograd computation.Comments