Pytorch相關（第一篇）

玥茹苟發表於2024-09-07
PyTorch

torch.autograd.Function 使用方法

torch.autograd.Function 是 PyTorch 提供的一個介面，用於自定義自動求導的操作。透過繼承這個類，你能夠定義自定義的前向和反向傳播邏輯。下面是使用 torch.autograd.Function 的基本步驟以及示例。

自定義 Function 的步驟

  1. 繼承 torch.autograd.Function。
  2. 實現 forward 和 backward 方法。
    • forward(ctx, *input)：計算前向傳播，並儲存需要在反向傳播中使用的任何資料。
    • backward(ctx, *grad_output)：計算反向傳播，根據輸出的梯度計算輸入的梯度。

示例程式碼

下面是一個簡單的例子，演示如何自定義一個函式，將輸入張量加倍：

import torch

class MyDoublingFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input):
        # 儲存中間結果到上下文
        ctx.save_for_backward(input)
        return input * 2  # 前向傳播：輸入加倍

    @staticmethod
    def backward(ctx, grad_output):
        # 從上下文中獲取輸入
        input, = ctx.saved_tensors
        # 反向傳播：梯度也加倍
        grad_input = grad_output.clone()  # 對 grad_output 進行克隆
        return grad_input  # 這裡實現了 d(output)/d(input)

# 使用自定義的 Function
x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
doubling = MyDoublingFunction.apply  # 獲取自定義函式的引用

# 計算前向傳播
y = doubling(x)
print("Output:", y)  # 輸出: tensor([2.0, 4.0, 6.0], grad_fn=<MyDoublingFunctionBackward>)

# 計算損失，反向傳播
loss = y.sum()
loss.backward()

# 列印梯度
print("Gradient:", x.grad)  # 輸出: tensor([2.0, 2.0, 2.0])

程式碼解釋

  1. 自定義類 MyDoublingFunction：

    • 繼承自 torch.autograd.Function。
    • 實現了 forward 和 backward 方法。
    • ctx.save_for_backward(input) 用於儲存輸入張量，以便在反向傳播中使用。

  2. 前向傳播：

    • 在 forward 方法中，輸入張量乘以 2，並返回結果。

  3. 反向傳播：

    • 在 backward 方法中，從上下文中獲取輸入，返回與 grad_output 相同的梯度，這樣在前向傳播中加倍的效果在反向傳播時也得到了相應的保持。

  4. 使用自定義 Function：

    • 建立一個包含梯度計算的張量 x。
    • 呼叫自定義的 doubling 函式進行前向傳播。
    • 計算損失並透過呼叫 loss.backward() 執行反向傳播，計算 x 的梯度。

重要注意事項

  • 靜態方法：forward 和 backward 必須是靜態方法，因為它們不會依賴於類的例項。
  • 上下文儲存：所有中間計算的資料應使用 ctx.save_for_backward() 儲存，並在 backward 方法中透過 ctx.saved_tensors 訪問。
  • 梯度計算：在 forward 和 backward 中，必須謹慎處理梯度。這一過程與定義正確的數學操作相結合，以確保反向傳播的準確性。

透過自定義 torch.autograd.Function，您可以靈活地實現任何需要的操作，同時能夠充分利用 PyTorch 的自動求導機制。

相關文章