"""
pytorch的hook（钩子）函数
"""
import torch
import torch.nn as nn
from tools.common_tools import set_seed

set_seed(1)  # 设置随机种子


# ----------------------------------- 1 tensor hook 1 -----------------------------------
flag = 0
#flag = 1
if flag:

    # 定义一个计算图 y = (w+x)*(w+1)
    w = torch.tensor([1.], requires_grad=True)
    x = torch.tensor([2.], requires_grad=True)
    a = torch.add(w, x)
    b = torch.add(w, 1)
    y = torch.mul(a, b)

    a_grad = list() # 存储张量的梯度

    def grad_hook(grad): # 这个函数的目的是保存（钩取）梯度
        a_grad.append(grad)

    handle = a.register_hook(grad_hook) # 将hook函数注册到张量a上

    y.backward() # 执行反向传播，梯度会被hook函数记录

    # 查看梯度
    print("gradient:", w.grad, x.grad, a.grad, b.grad, y.grad) # 只有叶子节点梯度仍然存在
    print("a_grad[0]: ", a_grad[0]) # 我们发现释放掉的a的梯度被记录下来了
    handle.remove() # 将记录的梯度释放掉


# ----------------------------------- 2 tensor hook 2 -----------------------------------
flag = 0
#flag = 1
if flag:

    w = torch.tensor([1.], requires_grad=True)
    x = torch.tensor([2.], requires_grad=True)
    a = torch.add(w, x)
    b = torch.add(w, 1)
    y = torch.mul(a, b)

    a_grad = list()

    def grad_hook(grad):
        grad *= 2 # 10 = 5*2
        return grad*3 # 不带return则返回10，否则会覆盖掉10，返回30=5*2*3

    handle = w.register_hook(grad_hook)

    y.backward()

    # 查看梯度
    print("w.grad: ", w.grad)
    handle.remove()


# ----------------------------------- 3 Module.register_forward_hook and pre hook -----------------------------------
#flag = 0
flag = 1
if flag:

    class Net(nn.Module):
        def __init__(self):
            super(Net, self).__init__()
            self.conv1 = nn.Conv2d(1, 2, 3) # in_channel = 1, out_channel = 2, kernel_size = 3
            self.pool1 = nn.MaxPool2d(2, 2) # kernel_size = 2, stride = 2

        def forward(self, x):
            x = self.conv1(x)
            x = self.pool1(x)
            return x

    def forward_hook(module, data_input, data_output):
        fmap_block.append(data_output) # 记录 feature map
        input_block.append(data_input) # 记录网络层输入数据

    def forward_pre_hook(module, data_input):
        print("forward_pre_hook input:{}".format(data_input))

    def backward_hook(module, grad_input, grad_output):
        print("backward hook input:{}".format(grad_input))
        print("backward hook output:{}".format(grad_output))

    # 初始化网络
    net = Net()
    net.conv1.weight[0].detach().fill_(1)
    net.conv1.weight[1].detach().fill_(2)
    net.conv1.bias.data.detach().zero_()

    # 注册hook，一共有前向、前向之前、反向传播三种hook
    fmap_block = list()
    input_block = list()
    net.conv1.register_forward_hook(forward_hook) # 这三个函数接受的都是一个函数
    net.conv1.register_forward_pre_hook(forward_pre_hook)
    net.conv1.register_backward_hook(backward_hook)

    # inference
    fake_img = torch.ones((1, 1, 4, 4))   # batch size * channel * H * W
    
    # 前向传播时，会先调用 register_forward_pre_hook 中注册的函数
    # 然后调用 register_forward_hook 中注册的函数
    output = net(fake_img) 

    loss_fnc = nn.L1Loss()
    target = torch.randn_like(output) # 随机定义一个目标，用于计算loss
    loss = loss_fnc(target, output)
    loss.backward() # 反向传播时，会调用register_backward_hook 中注册的函数

    # 观察
    print("output shape: {}\noutput value: {}\n".format(output.shape, output))
    print("feature maps shape: {}\noutput value: {}\n".format(fmap_block[0].shape, fmap_block[0]))
    print("input shape: {}\ninput value: {}".format(input_block[0][0].shape, input_block[0]))






















