Custom autograd.Function: must it be static?

username: tom-programming
cooked: <p>Hello,</p>
<p>I noticed that a custom Function can be built in two different ways:</p>
<p>(1)</p>
<pre><code class="lang-auto">class Test(torch.autograd.Function):
    def __init__(self):
        super(Test,self).__init__()

    def forward(self, x1, x2):
        self.state = state(x1)
        return torch.arange(8)

    def backward(self, grad_out):
        grad_input = grad_out.clone()
        return torch.arange(10,18),torch.arange(20,28)

# then use function = Test()
</code></pre>
<p>or<br>
(2)</p>
<pre><code class="lang-auto">class Test(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x1, x2):
        ctx.state = state(x1)
        return torch.arange(8)

    @staticmethod
    def backward(ctx, grad_out):
        grad_input = grad_out.clone()
        return torch.arange(10,18),torch.arange(20,28)

# then use function = Test.apply
</code></pre>
<p>The second option is the only to be referred in the official documentation, but it can be burdensome because if there are options as to how the function should be implemented, they must be initialized globally.</p>
<p>I prefer option (1) since I can save an intermediate variable (state) by myself and also initialize the function with some options, but it is not present in the official documentation</p>
<p>Am I safe in using option (1)?</p>

username: albanD
cooked: <p>Hi,</p>
<p>Option (1) is the old way to define <code>Function</code>s. This does not support gradients of gradients and it’s support might be discontinued in the future (not sure when).<br>
The second one is the way to go. Note that you can do exactly the same thing as you can save arbitrary stuff in the <code>ctx</code> (the same way you would save in <code>self</code> in (1)), and the <code>apply</code> method that calls <code>forward</code> accept any parameter, so you can just pass what you used to give the <code>__init__()</code> function here. That means that you don’t need to define options globally, just pass them to the forward method.</p>

username: Megh_Bhalerao
cooked: <p>Hi,<br>
Could someone please help me in converting this snippet to the newer way of defining custom layers:</p>
<pre><code class="lang-auto">class GradReverse(Function):
    def __init__(self, lambd):
        self.lambd = lambd
    #@staticmethod
    def forward(self, x):
        return x.view_as(x)
    #@staticmethod
    def backward(self, grad_output):
        return (grad_output * -self.lambd)```

</code></pre>
<p>Thanks a lot,<br>
Megh</p>

username: albanD
cooked: <p>Hi,</p>
<p>Sure! here you go:</p>
<pre><code class="lang-auto">class GradReverse(Function):
    @staticmethod
    def forward(ctx, x, lambd):
        ctx.lambd = lambd
        return x.view_as(x)

    @staticmethod
    def backward(ctx, grad_output):
        return (grad_output * -ctx.lambd), None
</code></pre>
<p>And you replace:</p>
<pre><code class="lang-auto">GradReverse(lambd)(inp)
</code></pre>
<p>with</p>
<pre><code class="lang-auto">GradReverse.apply(inp, lambd)
</code></pre>

username: Megh_Bhalerao
cooked: <p>Hi <a class="mention" href="/u/alband">@albanD</a>,<br>
Thanks a lot for your help!<br>
Megh</p>

username: QuanNguyen
cooked: <p>Hello <a class="mention" href="/u/alband">@albanD</a>, I have converted the custom autograd.Function code as new way, but I can’t save the torch model, do you have any suggests pls ?</p>
<pre><code class="lang-auto">class bin(torch.autograd.Function):
    """bin {0, 1} a real valued tensor."""

    @staticmethod
    def forward(ctx, inputs, threshold=DEFAULT_THRESHOLD):
        ctx.threshold = threshold
        outputs = inputs.clone()
        outputs[inputs.le(ctx.threshold)] = 0
        outputs[inputs.gt(ctx.threshold)] = 1
        return outputs

    @staticmethod
    def backward(ctx, gradOutput):
        return gradOutput, None
</code></pre>
<blockquote>
<p>torch.save(ckpt, savename)<br>
File “/home/quannguyen/.local/lib/python3.6/site-packages/torch/serialization.py”, line 372, in save<br>
_save(obj, opened_zipfile, pickle_module, pickle_protocol)<br>
File “/home/quannguyen/.local/lib/python3.6/site-packages/torch/serialization.py”, line 476, in _save<br>
pickler.dump(obj)<br>
TypeError: can’t pickle bin objects</p>
</blockquote>

username: _hu
cooked: <p>hello,i wanna ask the problem i have faced.I hope that you can help me.</p>
<pre data-code-wrap="python"><code class="lang-python">def conv_offset2d(input,
                  offset,
                  weight,
                  stride=1,
                  padding=0,
                  dilation=1,
                  deform_groups=1):

    if input is not None and input.dim() != 4:
        raise ValueError(
            "Expected 4D tensor as input, got {}D tensor instead.".format(
                input.dim()))
    stride=parse(stride)
    padding = parse(padding)
    dilation = parse(dilation)
    f = ConvOffset2dFunction(
        _pair(stride), _pair(padding), _pair(dilation), deform_groups)
    return f(input, offset, weight)


class ConvOffset2dFunction(Function):
    def __init__(self, stride, padding, dilation, deformable_groups=1):
        super(ConvOffset2dFunction, self).__init__()
        self.stride = stride
        self.padding = padding
        self.dilation = dilation
        self.deformable_groups = deformable_groups

    @staticmethod
    def forward(self, input, offset, weight):
        self.save_for_backward(input, offset, weight)

        output = input.new(*self._output_size(input, weight))

        self.bufs_ = [input.new(), input.new()]  # columns, ones

        if not input.is_cuda:
            raise NotImplementedError
        else:
            if isinstance(input, torch.autograd.Variable):
                if not isinstance(input.data, torch.cuda.FloatTensor):
                    raise NotImplementedError
            else:
                if not isinstance(input, torch.cuda.FloatTensor):
                    raise NotImplementedError
            deform_conv.deform_conv_forward_cuda(
                input, weight, offset, output, self.bufs_[0], self.bufs_[1],
                weight.size(3), weight.size(2), self.stride[1], self.stride[0],
                self.padding[1], self.padding[0], self.dilation[1],
                self.dilation[0], self.deformable_groups)
        return output

    @staticmethod
    def backward(self, grad_output):
        input, offset, weight = self.saved_tensors

        grad_input = grad_offset = grad_weight = None

        if not grad_output.is_cuda:
            raise NotImplementedError
        else:
            if isinstance(grad_output, torch.autograd.Variable):
                if not isinstance(grad_output.data, torch.cuda.FloatTensor):
                    raise NotImplementedError
            else:
                if not isinstance(grad_output, torch.cuda.FloatTensor):
                    raise NotImplementedError
            if self.needs_input_grad[0] or self.needs_input_grad[1]:
                grad_input = input.new(*input.size()).zero_()
                grad_offset = offset.new(*offset.size()).zero_()
                deform_conv.deform_conv_backward_input_cuda(
                    input, offset, grad_output, grad_input,
                    grad_offset, weight, self.bufs_[0], weight.size(3),
                    weight.size(2), self.stride[1], self.stride[0],
                    self.padding[1], self.padding[0], self.dilation[1],
                    self.dilation[0], self.deformable_groups)

            if self.needs_input_grad[2]:
                grad_weight = weight.new(*weight.size()).zero_()
                deform_conv.deform_conv_backward_parameters_cuda(
                    input, offset, grad_output,
                    grad_weight, self.bufs_[0], self.bufs_[1], weight.size(3),
                    weight.size(2), self.stride[1], self.stride[0],
                    self.padding[1], self.padding[0], self.dilation[1],
                    self.dilation[0], self.deformable_groups, 1)

        return grad_input, grad_offset, grad_weight

    def _output_size(self, input, weight):
        channels = weight.size(0)

        output_size = (input.size(0), channels)
        for d in range(input.dim() - 2):
            in_size = input.size(d + 2)
            pad = self.padding[d]
            kernel = self.dilation[d] * (weight.size(d + 2) - 1) + 1
            stride = self.stride[d]
            output_size += ((in_size + (2 * pad) - kernel) // stride + 1, )
        if not all(map(lambda s: s &gt; 0, output_size)):
            raise ValueError(
                "convolution input is too small (output would be {})".format(
                    'x'.join(map(str, output_size))))
        return output_size
</code></pre>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/7/f/7fc591f69516e0852ceb2d77cbc2a9196bed8390.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/7fc591f69516e0852ceb2d77cbc2a9196bed8390" title="捕获"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/7/f/7fc591f69516e0852ceb2d77cbc2a9196bed8390_2_690x87.jpeg" alt="捕获" data-base62-sha1="iejVIj7uRJFKWyfBuhxnoyxVe3m" width="690" height="87" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/7/f/7fc591f69516e0852ceb2d77cbc2a9196bed8390_2_690x87.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/7/f/7fc591f69516e0852ceb2d77cbc2a9196bed8390_2_1035x130.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/7/f/7fc591f69516e0852ceb2d77cbc2a9196bed8390_2_1380x174.jpeg 2x" data-dominant-color="F4F5F6"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">捕获</span><span class="informations">1610×203 51.8 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

username: albanD
cooked: <p><a class="mention" href="/u/_hu">@_hu</a> please follow the doc on how to write a custom Function <a href="https://pytorch.org/docs/stable/notes/extending.html">here</a> to avoid this error. <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>