Как именно работают прямые и обратные перехваты в PyTorch ⇐ Python

[Anonymous]

Я пытаюсь понять, как именно с точки зрения кода работают перехватчики в PyTorch. У меня есть модель, и я хотел бы установить в своем коде перехват вперед и назад. Я хотел бы установить привязку в моей модели после определенного слоя, и я думаю, что самый простой способ — установить привязку к этому конкретному модулю. Это вступительное видео предупреждает, что обратный модуль содержит ошибку, но я не уверен, что это все еще так.
Мой код выглядит следующим образом:

Код: Выделить всё

def __init__(self, model, attention_layer_name='desired_name_module',discard_ratio=0.9):
self.model = model
self.discard_ratio = discard_ratio
for name, module in self.model.named_modules():
if attention_layer_name in name:
module.register_forward_hook(self.get_attention)
module.register_backward_hook(self.get_attention_gradient)

self.attentions = []
self.attention_gradients = []

def get_attention(self, module, input, output):
self.attentions.append(output.cpu())

def get_attention_gradient(self, module, grad_input, grad_output):
self.attention_gradients.append(grad_input[0].cpu())

def __call__(self, input_tensor, category_index):
self.model.zero_grad()
output = self.model(input_tensor)
loss = ...
loss.backward()

Мне неясно, как с точки зрения кода работают следующие строки:

Код: Выделить всё

module.register_forward_hook(self.get_attention)
module.register_backward_hook(self.get_attention_gradient)

Я регистрирую привязку к нужному модулю, однако затем в каждом случае вызываю функцию без каких-либо входных данных. Мой вопрос касается Python: как именно работает этот вызов? Как аргументы Register_forward_hook и Register_backward_hook действуют при вызове функции?

Подробнее здесь: https://stackoverflow.com/questions/782 ... in-pytorch