#python #pytorch #tensor
Вопрос:
Я ищу хороший способ переопределения обратной операции в nn.Module, например:
class LayerWithCustomGrad(nn.Module):
def __init__(self):
super(LayerWithCustomGrad, self).__init__()
self.weights = nn.Parameter(torch.randn(200))
def forward(self,x):
return x * self.weights
def backward(self,grad_of_c): # This gets called during loss.backward()
# grad_of_c comes from the gradient of b*23
grad_of_a = some_operation(grad_of_c)
# perform extra computation
# and more computation
self.weights.grad = another_operation(grad_of_a,grad_of_c)
return grad_of_a # and the grad of parameter "a" will receive this
layer = LayerWithCustomGrad()
a = nn.Parameter(torch.randn(200),requires_grad=True)
b = layer(a)
c = b*23
Некоторые из проектов, над которыми я работаю, содержат слои с недифференцируемыми функциями, мне понравится, если есть способ соединить два разбитых графика и/или изменить градиенты уже существующих графиков.
Также будет здорово, если есть возможный способ сделать это в тензорном потоке
Комментарии:
1. Похоже, вы придерживаетесь правильного подхода. В чем именно заключается проблема?
2. Таким образом, в основном с помощью forward вы выполняете x -> операцию ->> больше операций ->>> результат. с обратным его grad_of_result -> к grad_of_more_operation и так далее. Цель состоит в том, чтобы повысить уровень grad_of_more_operation и изменить его до того, как он перейдет к «работе» во время потери.назад()
Ответ №1:
Способ, которым построен PyTorch, вы должны сначала реализовать пользовательский torch.autograd.Function
интерфейс, который будет содержать прямой и обратный проход для вашего слоя. Затем вы можете создать nn.Module
оболочку для этой функции с необходимыми параметрами.
На этой странице руководства вы можете увидеть, как реализуется ReLU. Я покажу здесь, как построить а torch.autograd.Function
и его nn.Module
обертку.
class F(torch.autograd.Function):
"""Both forward and backward are static methods."""
@staticmethod
def forward(ctx, input, weights):
"""
In the forward pass we receive a Tensor containing the input and return
a Tensor containing the output. ctx is a context object that can be used
to stash information for backward computation. You can cache arbitrary
objects for use in the backward pass using the ctx.save_for_backward method.
"""
ctx.save_for_backward(input, weights)
return input*weights
@staticmethod
def backward(ctx, grad_output):
"""
In the backward pass we receive a Tensor containing the gradient of the loss
with respect to the output, and we need to compute the gradient of the loss
with respect to the inputs: here input and weights
"""
input, weights = ctx.saved_tensors
grad_input = weights.clone()*grad_output
grad_weights = input.clone()*grad_output
return grad_input, grad_weights
Он nn.Module
инициализирует параметры и вызовет F
обработку фактического вычисления операции для прямого/обратного прохода.
class LayerWithCustomGrad(nn.Module):
def __init__(self):
super().__init__()
self.weights = nn.Parameter(torch.rand(10))
self.fn = F.apply
def forward(self, x):
return self.fn(x, self.weights)
Теперь мы можем попытаться сделать вывод и вернуться назад:
>>> layer = LayerWithCustomGrad()
>>> x = torch.randn(10, requires_grad=True)
>>> y = layer(x)
tensor([ 0.2023, 0.7176, 0.3577, -1.3573, 1.5185, 0.0632, 0.1210, 0.1566,
0.0709, -0.4324], grad_fn=<FBackward>)
Обратите внимание на <FBackward>
as grad_fn
: это обратная функция F
привязки к предыдущему выводу, с которым мы сделали x
.
>>> y.mean().backward()
>>> x.grad # i.e. grad_input in F.backward
tensor([0.0141, 0.0852, 0.0450, 0.0922, 0.0400, 0.0988, 0.0762, 0.0227, 0.0569,
0.0309])
>>> layer.weights.grad # i.e. grad_weights in F.backward
tensor([-1.4584, -2.1187, 1.5991, 0.9764, 1.8956, -1.0993, -3.7835, -0.4926,
0.9477, -1.2219])