How to get ctx.xx outside the custom autograd function?

Naruto-Sasuke · December 16, 2018, 7:57am

When I implement custom autograd function.
In the forward of class xxxFunction(torch.autograd.Function),
I get ctx.matches, then I want to get the ctx.matches outside this function.
I tried

@staticmethod
def get_matches(ctx):
    return ctx.matches

Then in the class xxx(torch.nn.Module), I run xxxFunction.get_matches().
get_matches() takes exactly 1 argument (0 given).

Thank you in advance!

Naruto-Sasuke · December 17, 2018, 1:47am

@albanD could you help me in this problem?

albanD · December 17, 2018, 10:37am

Hi,

You can access this from the backward method that gets the same ctx as input.
I don’t think you can extract it in a reliable way though.
Why do you need to do this? Can’t you just return it as another output of the forward method?

cosmozhang1988 · September 10, 2019, 2:37pm

I am facing the same issue. If I return it as another output of the forward method, how to adjust the backward method? @albanD
Many thanks!

albanD · September 10, 2019, 8:28pm

The backward method should get and return None for everything that is not differentiable.

Zhaoyi-Yan · September 17, 2019, 3:02am

Here is an example:

github.com

Zhaoyi-Yan/Shift-Net_pytorch/blob/master/models/shift_net/InnerShiftTriple.py#L34


    cur_device = torch.cuda.current_device()
    self.cur_mask = self.mask_all[cur_device*cur_bsize:(cur_device+1)*cur_bsize, :, :, :]


# If mask changes, then need to set cal_fix_flag true each iteration.
def forward(self, input):
    self.bz, self.c, self.h, self.w = input.size()
    self._split_mask(self.bz)
    self.flag = util.cal_flag_given_mask_thred(self.cur_mask, self.shift_sz, self.stride, self.mask_thred)
    final_out = InnerShiftTripleFunction.apply(input, self.shift_sz, self.stride, self.triple_weight, self.flag, self.show_flow)
    if self.show_flow:
        self.flow_srcs = InnerShiftTripleFunction.get_flow_src()
    return final_out


def get_flow(self):
    return self.flow_srcs


def set_flow_true(self):
    self.show_flow = True


def set_flow_false(self):
    self.show_flow = False