PyTorch no_grad

在進行 back-propagation 時，有些路徑是我們不想要去計算 gradient 的，這時我們可以用 no_grad() 這個 function。但有一個需要注意的地方是，只有在 no_grad() 底下產生出來的 tensor 才會被 disable gradient。以下舉兩個例子：

a = torch.tensor(1.0, requires_grad=True)
with torch.no_grad():
     b = a
b.backward()

這時候我們去看 a.grad，會發現還是有 gradient 產生。原因就在於，這個 b = a 其實只是讓 b 成為一個 a 的 reference，也就是 b 其實就是 a，只是名字換了而已。在這種情況下，b 不是一個在 no_grad() 底下產生的 tensor，因此還是會有 gradient。

a = torch.tensor(1.0, requires_grad=True)
with torch.no_grad():
     b = a + 0
b.backward()

假如改成 b = a + 0，那所產生的 b 就是一個新的 tensor，並且此 tensor 是不會被計算 gradient 的。

PyTorch 如何追蹤 backward 路徑

next_functions

當我們在建立一個 model 並進行訓練時，有時候可能 forward 結果是對的，但是訓練結果卻不如預期，這時候可能會想要知道 back propagation 的運算是否正確。這時候應該怎麼做呢？我們舉以下這段程式為例：

a = torch.tensor(1.1, requires_grad=True)
b = torch.tensor(1.2, requires_grad=True)
c = torch.tensor(1.3, requires_grad=True)
d = a * b
e = d * c
e.backward()

我們可以使用 e.grad_fn.next_functions 來追蹤 backward 的路徑。e.grad_fn.next_functions 是一個 list，其中包含了 e 往前走一層的所有 gradient function。上面這個例子中，e.grad_fn.next_functions 中包含兩個 element，一個是往 d 那邊走的 gradient function，另一個是往 c 那邊走的。

往 d 走的是 e.grad_fn.next_functions[0][0]
往 c 走的是 e.grad_fn.next_functions[1][0]

若我們想要知道 c 這條路徑的 forward 運算結果和 gradient，分別是：
Forward 結果：e.grad_fn.next_functions[1][0].variable   (結果為 1.3)
Gradient 大小：e.grad_fn.next_functions[1][0].variable.grad   (結果為 1.32)

而由於 d 這條路徑還有上游，我們必須再往上一層才能抵達 a 和 b

其中往 a 這條路徑的 forward 運算結果和 gradient，分別是：
Forward 結果：e.grad_fn.next_functions[0][0].next_functions[0][0].variable   (結果為 1.1)
Gradient 大小：e.grad_fn.next_functions[0][0].next_functions[0][0].variable.grad   (結果為 1.56 = 1.3 * 1.2)

而往 b 這條路徑的 forward 運算結果和 gradient，分別是：
Forward 結果：e.grad_fn.next_functions[0][0].next_functions[1][0].variable   (結果為 1.2)
Gradient 大小：e.grad_fn.next_functions[0][0].next_functions[1][0].variable.grad   (結果為 1.43 = 1.3 * 1.1)

只有位於 leaf node 的 grad_fn 有 variable 這個 attribute，因此必須使用 next_function 不斷的往前追蹤到所有的 leaf node 位置，過程中我們可以用 hasattr(object, name) 來確認 grad_fn 是否包含 variable 這個 attribute。

register_hook

如上所述，只有 graph 的 leaf node 存在 variable 能夠用於追蹤 gradient back-propagation。假如我們想要知道 internal node d 的 gradient 時，可以使用 register_hook。

承接上面這個例子，我們可以使用如下程式碼來取得 d 的 gradient。

grads = {}
def save_grad(name):
    def hook(grad):
        grads[name] = grad
    return hook
d.register_hook(save_grad('d'))

在做完 e.backward() 以後，就能夠用 grads['d'] 來存取 d 路徑上的 gradient，其結果為 1.3。

Reference
[1] https://stackoverflow.com/questions/52988876/how-can-i-visualize-what-happens-during-loss-backward
[2] https://codertw.com/%E7%A8%8B%E5%BC%8F%E8%AA%9E%E8%A8%80/368709/
[3] https://discuss.pytorch.org/t/why-cant-i-see-grad-of-an-intermediate-variable/94/6

PyTorch register_buffer

通常我們想要建立一個 tensor 時，會用以下方法：

a = torch.tensor(0.1)

假如要在一個 module 中建立一些 tensor，可能會使用如下方法：

class Layer(nn.Module):
    self.__init__(self, num_outputs, num_inputs):
        self.w = nn.Parameter(torch.Tensor(num_outputs, num_inputs))
        self.a = torch.tensor(0.1)

其中的 w 是可以被訓練的參數，而 a 是無法被訓練的。

這時我們會遇到一個問題，當我們想要將整個 module 從 CPU 移動到 GPU 時：

layer = Layer(10, 100)
layer = layer.to('cuda:0')

我們會發現 w 可以成功的移動，但是 a 卻仍然停留在 CPU 上。其中的差異就在，nn.Parameter 會將 w 宣告成一個 module 內的參數，但是用一般 torch.tensor 方法宣告的卻不會。

那麼假如我們想要在 module 內建立一個不可訓練的參數要怎麼做呢？答案是：使用 register_buffer 功能。

以上面的例子來說，a 可以改用以下方式宣告：

self.__init__(self):
    self.register_buffer('a', torch.tensor(0.1))

之後我們要使用 a 時，一樣用 self.a 就可以使用了。如此一來，當我們將 module 移動到 GPU 時，a 也會跟著一起移動。