Do I need some decorator for cpp operators when I use pytorch1.6?

coincheung: <p>Hi,</p>
<p>I am trying to implement an operator that will change the labels for computing the loss function. The loss function is like this:</p>
<pre><code class="lang-auto">class OhemCELoss(nn.Module):

    def __init__(self, thresh, gnore_lb=255):
        super(OhemCELoss, self).__init__()
        self.score_thresh = thresh
        self.ignore_lb = ignore_lb
        self.criteria = nn.CrossEntropyLoss(ignore_index=ignore_lb, reduction='mean')
    def forward(self, logits, labels):
        import ohem_cpp
        n_min = (labels != self.ignore_lb).numel() // 16
        labels = ohem_cpp.score_ohem_label(logits, labels,
                self.ignore_lb, self.score_thresh, n_min).detach()
        loss = self.criteria(logits, labels)
        return loss
</code></pre>
<p>Here the <code>ohem_cpp</code> is implemented according to the principle of the pytorch extension, and the <code>ohem_cpp.score_ohem_label</code> will compute the softmax scores of the input logits, choose the scores lower than the <code>thresh</code> and set the corresponding element of the <code>labels</code> to the assigned <code>ignore_lb</code>. I noticed that in pytorch1.6, users should add decorator of <code>@custum_fwd</code> if they want to implement their own operators that is derived from <code>autograd.Function</code>. I have no idea whether I can call my <code>ohem_cpp</code> module directly as in the code above without potential errors. Did I miss anything here ?</p>
<p>I am asking this because I sometimes met the error of <code>terminate called after throwing an instance of 'thrust::system::system_error'</code>(not often), I do not know whether is is associated with my wrongly using of my self-implemented operators.</p>

albanD: <aside class="quote no-group" data-username="coincheung" data-post="1" data-topic="92498">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/c/ad7895/48.png" class="avatar"> coincheung:</div>
<blockquote>
<p>I noticed that in pytorch1.6, users should add decorator of <code>@custum_fwd</code> if they want to implement their own operators that is derived from <code>autograd.Function</code></p>
</blockquote>
</aside>
<p>Where did you saw that? There were no changes that I’m aware of to <code>autograd.Function</code> in 1.6.</p>
<blockquote>
<p>I have no idea whether I can call my <code>ohem_cpp</code> module directly as in the code above without potential errors.</p>
</blockquote>
<p>Yes you can. And as long as you only use differentiable ops on the c++ side you can backprop through it.<br>
If you don’t, then you’ll need a custom Function and you can see <a href="https://pytorch.org/docs/stable/notes/extending.html">here</a> how to do it.</p>

coincheung: <p>Thanks, I got the <code>@custum_fwd</code> specification from the <a href="https://pytorch.org/docs/stable/notes/amp_examples.html#functions-with-multiple-inputs-or-autocastable-ops" rel="noopener nofollow ugc">amp tutorial</a>.</p>
<blockquote>
<p>Apply <a href="https://pytorch.org/docs/stable/amp.html#torch.cuda.amp.custom_fwd" rel="noopener nofollow ugc"> <code>custom_fwd</code> </a> and <a href="https://pytorch.org/docs/stable/amp.html#torch.cuda.amp.custom_bwd" rel="noopener nofollow ugc"> <code>custom_bwd</code> </a> (with no arguments) to <code>forward</code> and <code>backward</code> respectively.</p>
</blockquote>

albanD: <p>Thanks for the link. I wasn’t aware of this one <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=9" title=":smiley:" class="emoji" alt=":smiley:"></p>
<p>It looks orthogonal to the choice of using a custom Function or not.<br>
As mentioned above, if you c++ code is differentiable, you’re all good. If not, you’ll have to define the backward and can do so with a custom Function in python as instructed above that just happens to call you c++ code.</p>