目錄
快捷方式
intermediate/parametrizations

注意

點擊 這裡 下載完整的範例程式碼

參數化教學

建立於：2021 年 4 月 19 日 | 最後更新：2024 年 2 月 05 日 | 最後驗證：2024 年 11 月 05 日

作者: Mario Lezcano

正則化深度學習模型是一項非常具有挑戰性的任務。當應用於深度模型時，由於被優化的函數的複雜性，諸如懲罰方法之類的經典技術通常會失效。當使用病態模型時，這尤其成問題。這些模型的例子包括在長序列上訓練的 RNN 和 GAN。近年來，已經提出了許多技術來正則化這些模型並改善它們的收斂性。在循環模型上，已經提出控制 RNN 循環核心的奇異值以使其具有良好的條件。例如，可以通過使循環核心正交來實現。另一種正則化循環模型的方法是通過“權重歸一化”。這種方法建議將參數的學習與其範數的學習分開。為此，參數除以其Frobenius 範數，並且學習一個單獨的參數來編碼其範數。類似的正則化方法以“譜歸一化”的名稱提出用於 GAN。此方法通過將其參數除以其譜範數而不是其 Frobenius 範數來控制網路的 Lipschitz 常數。

所有這些方法都有一個共同的模式：它們都在使用參數之前以適當的方式轉換參數。在第一種情況下，它們使用將矩陣映射到正交矩陣的函數使其正交。在權重和譜歸一化的情況下，它們將原始參數除以其範數。

更一般地說，所有這些例子都使用函數在參數上放置額外的結構。換句話說，它們使用函數來約束參數。

在本教程中，您將學習如何實作和使用此模式來對您的模型施加約束。這樣做就像編寫您自己的 nn.Module 一樣容易。

要求：torch>=1.9.0

手動實作參數化

假設我們想要一個具有對稱權重的方形線性層，也就是說，權重 X 使得 X = Xᵀ。一種方法是將矩陣的上三角部分複製到其下三角部分

import torch
import torch.nn as nn
import torch.nn.utils.parametrize as parametrize

def symmetric(X):
    return X.triu() + X.triu(1).transpose(-1, -2)

X = torch.rand(3, 3)
A = symmetric(X)
assert torch.allclose(A, A.T)  # A is symmetric
print(A)                       # Quick visual check

tensor([[0.8823, 0.9150, 0.3829],
        [0.9150, 0.3904, 0.6009],
        [0.3829, 0.6009, 0.9408]])

然後，我們可以利用這個想法來實作一個具有對稱權重的線性層

class LinearSymmetric(nn.Module):
    def __init__(self, n_features):
        super().__init__()
        self.weight = nn.Parameter(torch.rand(n_features, n_features))

    def forward(self, x):
        A = symmetric(self.weight)
        return x @ A

然後，該層可以像常規線性層一樣使用

layer = LinearSymmetric(3)
out = layer(torch.rand(8, 3))

這個實作雖然正確且獨立，但存在許多問題

  1. 它重新實作了該層。我們必須將線性層實作為 x @ A。對於線性層來說，這不是很有問題，但想像一下必須重新實作 CNN 或 Transformer……

  2. 它沒有將層和參數化分開。如果參數化更困難，我們必須為我們想要在其中使用的每個層重寫其程式碼。

  3. 每次使用該層時，它都會重新計算參數化。如果我們在前向傳遞過程中多次使用該層（想像一下 RNN 的循環核心），它會在每次呼叫該層時計算相同的 A。

參數化簡介

參數化可以解決所有這些問題以及其他問題。

讓我們首先使用 torch.nn.utils.parametrize 重新實作上面的程式碼。我們唯一需要做的是將參數化編寫為常規 nn.Module

class Symmetric(nn.Module):
    def forward(self, X):
        return X.triu() + X.triu(1).transpose(-1, -2)

這就是我們需要做的全部。一旦我們有了這個，我們可以通過執行以下操作將任何常規層轉換為對稱層

layer = nn.Linear(3, 3)
parametrize.register_parametrization(layer, "weight", Symmetric())

ParametrizedLinear(
  in_features=3, out_features=3, bias=True
  (parametrizations): ModuleDict(
    (weight): ParametrizationList(
      (0): Symmetric()
    )
  )
)

現在，線性層的矩陣是對稱的

A = layer.weight
assert torch.allclose(A, A.T)  # A is symmetric
print(A)                       # Quick visual check

tensor([[ 0.2430,  0.5155,  0.3337],
        [ 0.5155,  0.3333,  0.1033],
        [ 0.3337,  0.1033, -0.5715]], grad_fn=<AddBackward0>)

我們可以對任何其他層執行相同的操作。例如，我們可以創建一個具有 斜對稱 核心的 CNN。我們使用類似的參數化，將符號反轉的上三角部分複製到下三角部分

class Skew(nn.Module):
    def forward(self, X):
        A = X.triu(1)
        return A - A.transpose(-1, -2)


cnn = nn.Conv2d(in_channels=5, out_channels=8, kernel_size=3)
parametrize.register_parametrization(cnn, "weight", Skew())
# Print a few kernels
print(cnn.weight[0, 1])
print(cnn.weight[2, 2])

tensor([[ 0.0000,  0.0457, -0.0311],
        [-0.0457,  0.0000, -0.0889],
        [ 0.0311,  0.0889,  0.0000]], grad_fn=<SelectBackward0>)
tensor([[ 0.0000, -0.1314,  0.0626],
        [ 0.1314,  0.0000,  0.1280],
        [-0.0626, -0.1280,  0.0000]], grad_fn=<SelectBackward0>)

檢查參數化模組

當模組被參數化時，我們發現模組以三種方式發生了變化

  1. model.weight 現在是一個屬性

  2. 它有一個新的 module.parametrizations 屬性

  3. 未參數化的權重已移動到 module.parametrizations.weight.original


在參數化 weight 之後，layer.weight 被轉換為 Python 屬性。每次我們請求 layer.weight 時，這個屬性都會計算 parametrization(weight)，就像我們在上面實作 LinearSymmetric 時所做的一樣。

已註冊的參數化會儲存在模組內的 parametrizations 屬性下。

layer = nn.Linear(3, 3)
print(f"Unparametrized:\n{layer}")
parametrize.register_parametrization(layer, "weight", Symmetric())
print(f"\nParametrized:\n{layer}")

Unparametrized:
Linear(in_features=3, out_features=3, bias=True)

Parametrized:
ParametrizedLinear(
  in_features=3, out_features=3, bias=True
  (parametrizations): ModuleDict(
    (weight): ParametrizationList(
      (0): Symmetric()
    )
  )
)

這個 parametrizations 屬性是一個 nn.ModuleDict，可以這樣存取它

print(layer.parametrizations)
print(layer.parametrizations.weight)

ModuleDict(
  (weight): ParametrizationList(
    (0): Symmetric()
  )
)
ParametrizationList(
  (0): Symmetric()
)

這個 nn.ModuleDict 的每個元素都是一個 ParametrizationList，它的行為類似於 nn.Sequential。這個列表允許我們將多個參數化串聯在同一個權重上。由於這是一個列表，我們可以透過索引來存取這些參數化。我們的 Symmetric 參數化就在這裡

print(layer.parametrizations.weight[0])

Symmetric()

我們注意到的另一件事是，如果我們印出這些參數，我們會看到參數 weight 已經被移動了

print(dict(layer.named_parameters()))

{'bias': Parameter containing:
tensor([-0.0730, -0.2283,  0.3217], requires_grad=True), 'parametrizations.weight.original': Parameter containing:
tensor([[-0.4328,  0.3425,  0.4643],
        [ 0.0937, -0.1005, -0.5348],
        [-0.2103,  0.1470,  0.2722]], requires_grad=True)}

它現在位於 layer.parametrizations.weight.original 下

print(layer.parametrizations.weight.original)

Parameter containing:
tensor([[-0.4328,  0.3425,  0.4643],
        [ 0.0937, -0.1005, -0.5348],
        [-0.2103,  0.1470,  0.2722]], requires_grad=True)

除了這三個小差異之外，這個參數化的行為與我們手動實作的完全相同

symmetric = Symmetric()
weight_orig = layer.parametrizations.weight.original
print(torch.dist(layer.weight, symmetric(weight_orig)))

tensor(0., grad_fn=<DistBackward0>)

參數化是一等公民

由於 layer.parametrizations 是一個 nn.ModuleList，這表示這些參數化已正確地註冊為原始模組的子模組。因此，在模組中註冊參數的相同規則也適用於註冊參數化。例如，如果一個參數化有參數，當呼叫 model = model.cuda() 時，這些參數將會從 CPU 移至 CUDA。

快取參數化的值

參數化帶有一個內建的快取系統，透過 context manager parametrize.cached() 來實現

class NoisyParametrization(nn.Module):
    def forward(self, X):
        print("Computing the Parametrization")
        return X

layer = nn.Linear(4, 4)
parametrize.register_parametrization(layer, "weight", NoisyParametrization())
print("Here, layer.weight is recomputed every time we call it")
foo = layer.weight + layer.weight.T
bar = layer.weight.sum()
with parametrize.cached():
    print("Here, it is computed just the first time layer.weight is called")
    foo = layer.weight + layer.weight.T
    bar = layer.weight.sum()

Computing the Parametrization
Here, layer.weight is recomputed every time we call it
Computing the Parametrization
Computing the Parametrization
Computing the Parametrization
Here, it is computed just the first time layer.weight is called
Computing the Parametrization

串聯參數化

串聯兩個參數化就像在同一個 tensor 上註冊它們一樣容易。我們可以利用這個方法，從較簡單的參數化建立更複雜的參數化。例如，Cayley 映射將斜對稱矩陣映射到正定正交矩陣。我們可以串聯 Skew 和一個實作 Cayley 映射的參數化，以獲得一個具有正交權重的層

class CayleyMap(nn.Module):
    def __init__(self, n):
        super().__init__()
        self.register_buffer("Id", torch.eye(n))

    def forward(self, X):
        # (I + X)(I - X)^{-1}
        return torch.linalg.solve(self.Id - X, self.Id + X)

layer = nn.Linear(3, 3)
parametrize.register_parametrization(layer, "weight", Skew())
parametrize.register_parametrization(layer, "weight", CayleyMap(3))
X = layer.weight
print(torch.dist(X.T @ X, torch.eye(3)))  # X is orthogonal

tensor(2.8527e-07, grad_fn=<DistBackward0>)

這也可以用來修剪參數化的模組，或重複使用參數化。例如，矩陣指數將對稱矩陣映射到對稱正定 (SPD) 矩陣。但矩陣指數也將斜對稱矩陣映射到正交矩陣。利用這兩個事實，我們可以善用之前的參數化

class MatrixExponential(nn.Module):
    def forward(self, X):
        return torch.matrix_exp(X)

layer_orthogonal = nn.Linear(3, 3)
parametrize.register_parametrization(layer_orthogonal, "weight", Skew())
parametrize.register_parametrization(layer_orthogonal, "weight", MatrixExponential())
X = layer_orthogonal.weight
print(torch.dist(X.T @ X, torch.eye(3)))         # X is orthogonal

layer_spd = nn.Linear(3, 3)
parametrize.register_parametrization(layer_spd, "weight", Symmetric())
parametrize.register_parametrization(layer_spd, "weight", MatrixExponential())
X = layer_spd.weight
print(torch.dist(X, X.T))                        # X is symmetric
print((torch.linalg.eigvalsh(X) > 0.).all())  # X is positive definite

tensor(1.9066e-07, grad_fn=<DistBackward0>)
tensor(4.2147e-08, grad_fn=<DistBackward0>)
tensor(True)

初始化參數化

參數化帶有一個初始化它們的機制。如果我們實作一個簽名為 right_inverse 的方法

def right_inverse(self, X: Tensor) -> Tensor

它將會在指定給參數化的 tensor 時被使用。

讓我們升級我們的 Skew 類別的實作，以支援這個功能

class Skew(nn.Module):
    def forward(self, X):
        A = X.triu(1)
        return A - A.transpose(-1, -2)

    def right_inverse(self, A):
        # We assume that A is skew-symmetric
        # We take the upper-triangular elements, as these are those used in the forward
        return A.triu(1)

我們現在可以初始化一個用 Skew 參數化的層

layer = nn.Linear(3, 3)
parametrize.register_parametrization(layer, "weight", Skew())
X = torch.rand(3, 3)
X = X - X.T                             # X is now skew-symmetric
layer.weight = X                        # Initialize layer.weight to be X
print(torch.dist(layer.weight, X))      # layer.weight == X

tensor(0., grad_fn=<DistBackward0>)

當我們串聯參數化時，這個 right_inverse 會如預期般運作。為了看到這一點，讓我們升級 Cayley 參數化，使其也支援初始化

class CayleyMap(nn.Module):
    def __init__(self, n):
        super().__init__()
        self.register_buffer("Id", torch.eye(n))

    def forward(self, X):
        # Assume X skew-symmetric
        # (I + X)(I - X)^{-1}
        return torch.linalg.solve(self.Id - X, self.Id + X)

    def right_inverse(self, A):
        # Assume A orthogonal
        # See https://en.wikipedia.org/wiki/Cayley_transform#Matrix_map
        # (A - I)(A + I)^{-1}
        return torch.linalg.solve(A + self.Id, self.Id - A)

layer_orthogonal = nn.Linear(3, 3)
parametrize.register_parametrization(layer_orthogonal, "weight", Skew())
parametrize.register_parametrization(layer_orthogonal, "weight", CayleyMap(3))
# Sample an orthogonal matrix with positive determinant
X = torch.empty(3, 3)
nn.init.orthogonal_(X)
if X.det() < 0.:
    X[0].neg_()
layer_orthogonal.weight = X
print(torch.dist(layer_orthogonal.weight, X))  # layer_orthogonal.weight == X

tensor(2.2141, grad_fn=<DistBackward0>)

這個初始化步驟可以更簡潔地寫成

layer_orthogonal.weight = nn.init.orthogonal_(layer_orthogonal.weight)

這個方法的名稱來自於我們通常期望 forward(right_inverse(X)) == X。這是重寫在用值 X 初始化後，forward 應該傳回值 X 的直接方法。這個約束在實務中並未被嚴格執行。事實上，有時候，放寬這種關係可能是有意義的。例如，考慮以下隨機修剪方法的實作

class PruningParametrization(nn.Module):
    def __init__(self, X, p_drop=0.2):
        super().__init__()
        # sample zeros with probability p_drop
        mask = torch.full_like(X, 1.0 - p_drop)
        self.mask = torch.bernoulli(mask)

    def forward(self, X):
        return X * self.mask

    def right_inverse(self, A):
        return A

在這種情況下，對於每個矩陣 A，forward(right_inverse(A)) == A 並非總是成立。只有當矩陣 A 在與 mask 相同的位置有零時，這才成立。即使如此，如果我們將一個 tensor 指定給一個被修剪的參數，這個 tensor 事實上將會被修剪，這並不令人驚訝

layer = nn.Linear(3, 4)
X = torch.rand_like(layer.weight)
print(f"Initialization matrix:\n{X}")
parametrize.register_parametrization(layer, "weight", PruningParametrization(layer.weight))
layer.weight = X
print(f"\nInitialized weight:\n{layer.weight}")

Initialization matrix:
tensor([[0.3513, 0.3546, 0.7670],
        [0.2533, 0.2636, 0.8081],
        [0.0643, 0.5611, 0.9417],
        [0.5857, 0.6360, 0.2088]])

Initialized weight:
tensor([[0.3513, 0.3546, 0.7670],
        [0.2533, 0.0000, 0.8081],
        [0.0643, 0.5611, 0.9417],
        [0.5857, 0.6360, 0.0000]], grad_fn=<MulBackward0>)

移除參數化

我們可以透過使用 parametrize.remove_parametrizations() 來移除模組中參數或緩衝區的所有參數化

layer = nn.Linear(3, 3)
print("Before:")
print(layer)
print(layer.weight)
parametrize.register_parametrization(layer, "weight", Skew())
print("\nParametrized:")
print(layer)
print(layer.weight)
parametrize.remove_parametrizations(layer, "weight")
print("\nAfter. Weight has skew-symmetric values but it is unconstrained:")
print(layer)
print(layer.weight)

Before:
Linear(in_features=3, out_features=3, bias=True)
Parameter containing:
tensor([[ 0.0669, -0.3112,  0.3017],
        [-0.5464, -0.2233, -0.1125],
        [-0.4906, -0.3671, -0.0942]], requires_grad=True)

Parametrized:
ParametrizedLinear(
  in_features=3, out_features=3, bias=True
  (parametrizations): ModuleDict(
    (weight): ParametrizationList(
      (0): Skew()
    )
  )
)
tensor([[ 0.0000, -0.3112,  0.3017],
        [ 0.3112,  0.0000, -0.1125],
        [-0.3017,  0.1125,  0.0000]], grad_fn=<SubBackward0>)

After. Weight has skew-symmetric values but it is unconstrained:
Linear(in_features=3, out_features=3, bias=True)
Parameter containing:
tensor([[ 0.0000, -0.3112,  0.3017],
        [ 0.3112,  0.0000, -0.1125],
        [-0.3017,  0.1125,  0.0000]], requires_grad=True)

當移除一個參數化時，我們可以選擇保留原始參數（即 layer.parametrizations.weight.original 中的那個參數），而不是它的參數化版本，方法是設定 flag leave_parametrized=False

layer = nn.Linear(3, 3)
print("Before:")
print(layer)
print(layer.weight)
parametrize.register_parametrization(layer, "weight", Skew())
print("\nParametrized:")
print(layer)
print(layer.weight)
parametrize.remove_parametrizations(layer, "weight", leave_parametrized=False)
print("\nAfter. Same as Before:")
print(layer)
print(layer.weight)

Before:
Linear(in_features=3, out_features=3, bias=True)
Parameter containing:
tensor([[-0.3447, -0.3777,  0.5038],
        [ 0.2042,  0.0153,  0.0781],
        [-0.4640, -0.1928,  0.5558]], requires_grad=True)

Parametrized:
ParametrizedLinear(
  in_features=3, out_features=3, bias=True
  (parametrizations): ModuleDict(
    (weight): ParametrizationList(
      (0): Skew()
    )
  )
)
tensor([[ 0.0000, -0.3777,  0.5038],
        [ 0.3777,  0.0000,  0.0781],
        [-0.5038, -0.0781,  0.0000]], grad_fn=<SubBackward0>)

After. Same as Before:
Linear(in_features=3, out_features=3, bias=True)
Parameter containing:
tensor([[ 0.0000, -0.3777,  0.5038],
        [ 0.0000,  0.0000,  0.0781],
        [ 0.0000,  0.0000,  0.0000]], requires_grad=True)

指令碼的總執行時間： ( 0 分鐘 0.055 秒)

由 Sphinx-Gallery 產生的圖片集

文件

取得 PyTorch 的完整開發者文件

檢視文件

教學

取得針對初學者和進階開發者的深入教學

檢視教學

資源

尋找開發資源並獲得您的問題解答

檢視資源

為了分析流量並優化您的體驗，我們會在網站上使用 Cookie。 點擊或瀏覽即表示您同意允許我們使用 Cookie。 作為本網站目前的維護者，Facebook 的 Cookie 政策適用。 了解更多，包括可用的控制項：Cookie 政策。