torch::jit::script::Module with hook function?

Edwinzero: <p>Is there any way to compute the gradient of activation map at a certain later with respect to the one of the output components with torch::jit::script::Module?<br>
it seems the hook function only supports the nn::Module.<br>
Could you provide a simple example?</p>
<p>Thank you very much in advance.</p>

driazati: <p>Hooks aren’t supported right now but we are working on it and it will probably be done in the next couple weeks</p>

Edwinzero: <p>Thank you very much for your reply <a class="mention" href="/u/driazati">@driazati</a> !<br>
I am looking forward to that feature but however, I am catching a due with in one week, I wonder is there any alternate way to compute the gradient of hidden activation map at certain layer with jit::script::Module?</p>
<p>Thank you very much for your attention to this matter.</p>

driazati: <p>We don’t have any support for hooks so anything you do is going to be pretty hacky; that said, you might be able to get something working by adding a custom operator to TorchScript that just calls a C++ <code>autograd::Function</code> that does nothing except let’s you grab the gradients in the backward pass.</p>
<p>There’s an example of a custom op used in this way in torchvision, see <a href="https://github.com/pytorch/vision/blob/master/torchvision/csrc/empty_tensor_op.h">https://github.com/pytorch/vision/blob/master/torchvision/csrc/empty_tensor_op.h</a> and <a href="https://github.com/pytorch/vision/blob/43e94b39bcdda519c093ca11d99dfa2568aa7258/torchvision/csrc/vision.cpp#L51">https://github.com/pytorch/vision/blob/43e94b39bcdda519c093ca11d99dfa2568aa7258/torchvision/csrc/vision.cpp#L51</a></p>

Edwinzero: <aside class="quote no-group quote-modified" data-username="driazati" data-post="4" data-topic="75874">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/driazati/48/15393_2.png" class="avatar"> driazati:</div>
<blockquote>
<p>Hooks aren’t supported right now but we are working on it and it will probably be done in the next couple weeks</p>
</blockquote>
</aside>
<p>Hello <a class="mention" href="/u/driazati">@driazati</a>, does current LibTorch1.5 supports hook in the jit model?<br>
If so, could you please give an example of using hook function?<br>
Thank you very much for your attention to this matter.</p>

ChamSu: <p>Hi <a class="mention" href="/u/driazati">@driazati</a>, does current LibTorch1.7 supports hook in the jit model?<br>
I want to check whether the output of each layer is correct, and I found a forward_hook function play an important role in it.</p>

SplitInfinity: <p>1.7 does not but 1.8 does!</p>

AvalonChang: <p>Hello <a class="mention" href="/u/splitinfinity">@SplitInfinity</a>, do you have any example for jit::script::Module using hook function?<br>
I can’t found about hook function in LibTorch v1.10.2 (except at::Tensor::register_hook)</p>
<p>Thanks in advance.</p>

Tihaw: <p>I tried to save the outputs of hidden layer during forward, and then using autograd to calc the hidden layer’s gradients.<br>
Then it was good, the results are same as backward hooks do. But only can readout and calc the gradients, not modifying them.<br>
For my side, JIT is for inferring the model rather than training. Inferring stage needs some gradients from the hidden layer, and in this case, backward hooks is not necessary. <a href="https://github.com/pytorch/pytorch/issues/64728" rel="noopener nofollow ugc">jit autograd</a></p>