Dacă o operație matematică nu poate fi exprimată prin operatorii nativi PyTorch sau dacă dorim să implementăm un operator cu gradient trunchiat (Straight-Through Estimator – STE) pentru cuantizare, putem defini o sub-clasă a lui torch.autograd.Function.
import torch
class BinarizeSTE(torch.autograd.Function):
@staticmethod
def forward(ctx, x):
ctx.save_for_backward(x)
return torch.sign(x)
@staticmethod
def backward(ctx, grad_output):
x, = ctx.saved_tensors
# Straight-Through Estimator: gradientul trece neschimbat dacă |x| <= 1
grad_input = grad_output.clone()
grad_input[x.abs() > 1] = 0
return grad_input
# Utilizare în model
binarize = BinarizeSTE.apply
x = torch.tensor([0.5, -1.2, 0.8], requires_grad=True)
y = binarize(x)
y.sum().backward()
print("Gradiente calculate prin STE:", x.grad)Metoda ctx.save_for_backward stochează în siguranță tensorii necesari derivării în memoria tampon a contextului.
Întrebări & Exemple practice:
Extinderea PyTorch cu torch.autograd.Function PersonalizatăAi o întrebare despre Extinderea PyTorch cu torch.autograd.Function Personalizată sau ai descoperit un truc, prompt sau fragment de cod util? Distribuie-l cu ceilalți membri ai comunității!
Fii primul care pune o întrebare sau distribuie un fragment de cod / prompt!
Nu au fost adăugate încă întrebări sau resurse pentru această pagină. Dacă ai testat un prompt sau ai un snippet util, adaugă-l mai jos pentru comunitate.