2
$\begingroup$

I need to quantize the inputs, but the method (bucketize) I need to do so is indifferentiable. I can of course detach the tensor, but then I lose the flow of gradients to earlier weights. I guess the question is quite simple, how do you continue the flow of gradients when necessary. For example, using the following code ...

x = self.linear1(x)        
min, max = int(x.min()), int(x.max())        
bins = torch.linspace(min, max+1, 16)
x = torch.bucketize(x.detach(), bins) # forced to detach here
x = self.linear2(x)

I know it's possible, as it was done with VQVAE and the gradients still flow fine for their quantization purposes. But I checked several versions of VQVAE code and it doesn't seem logical how the gradients pass through the argmin method. It just seems to work regardless, which confuses me. I'm quite perplexed by this. I would be grateful for any help on this one. Thanks in advance.

$\endgroup$
1
  • 1
    $\begingroup$ Whenever you backpropagate, keep the retain_graph=True argument in the backward method. $\endgroup$
    – thanatoz
    Dec 30, 2020 at 6:42

1 Answer 1

1
$\begingroup$

Thank you @thanatoz for answering that in the comments. In fact, this did solve my problem. I also realized from other people's VQVAE code, there is also another possible solution in which the retrain_graph parameter is not needed. My example code could be written ...

x = self.linear1(x)        
min, max = int(x.min()), int(x.max())        
bins = torch.linspace(min, max+1, 16)
x_buckets = torch.bucketize(x.detach(), bins) # forced to detach here
x = x + (x - x_buckets).detach() # Reintroduce gradients here. <------------
x = self.linear2(x)

The answer is in the fifth line, by subtracting and then adding back again. Gradients can then be reintroduced.

$\endgroup$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge that you have read and understand our privacy policy and code of conduct.

Not the answer you're looking for? Browse other questions tagged or ask your own question.