Pytorch с модифицированными производными

Question

Я сейчас переписываю модель TRACX2 , разновидность рекуррентной нейронной сети, используемой для обучения кодировок в контексте сегментации слов из непрерывной речи или текста. Автор исходного кода вручную написал сеть в Numpy, а я хочу оптимизировать ее с помощью Pytorch. Однако они реализуют то, что они называют «температурой» и «смещением по Фалману»:

\ frac {d (tanh (x)))} {dx} = температура * (1 - tanh ^ 2 (x)))) + fahlmanOffset

Это явно не фактическая производная от tanh (x), одной из их функций активации, но вместо этого они использовали эту производную. Как я могу go реализовать эту модификацию в Pytorch?

Answer 1

По сути, вы добавляете обратный хук вот так:

a = Variable(torch.randn(2,2), requires_grad=True)
m = nn.Linear(2,1)
m(a).mean().backward()
print(a.grad) 
# shows a 2x2 tensor of non-zero values
temperature = 0.3
fahlmanOffset = .1

def hook(module, grad_input, grad_output):
    # Use custom gradient output
    return grad_output * temperature + fahlmanOffset

m.register_backward_hook(hook)

a.grad.zero_()
m(a).mean().backward()
print(a.grad)
# shows a 2x2 tensor with modified gradient

(любезно предоставлено этот ответ )