評價此頁
★ ★ ★ ★ ★
intermediate/parametrizations

注意

轉到末尾 下載完整的示例程式碼。

引數化教程#

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

作者： Mario Lezcano

對深度學習模型進行正則化是一項出人意料的挑戰性任務。由於被最佳化函式的複雜性，像罰函式法這樣的經典技術在應用於深度模型時往往力不從心。這在處理病態模型時尤其成問題。例如，在長序列上訓練的 RNN 和 GAN 就是這類模型。近年來，已提出了許多技術來正則化這些模型並改善它們的收斂性。對於迴圈模型，有人提出控制 RNN 的迴圈核的奇異值以使其保持良好病態。例如，可以透過使迴圈核 正交 來實現這一點。另一種正則化迴圈模型的方法是透過“權重歸一化”。這種方法建議將引數的學習與其範數鐒學習分離開來。為此，將引數除以其 Frobenius 範數，並學習一個單獨的引數來編碼其範數。GANs 中也提出了類似的正則化方法，稱為“譜歸一化”。該方法透過將引數除以它們的 譜範數（而不是 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.2994, 0.1108, 0.7758],
        [0.1108, 0.8337, 0.3246],
        [0.7758, 0.3246, 0.2519]])

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

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.0900,  0.2098,  0.2589],
        [ 0.2098, -0.3692, -0.5678],
        [ 0.2589, -0.5678,  0.0852]], 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.0053, -0.1311],
        [-0.0053,  0.0000,  0.1016],
        [ 0.1311, -0.1016,  0.0000]], grad_fn=<SelectBackward0>)
tensor([[ 0.0000,  0.0430, -0.1106],
        [-0.0430,  0.0000, -0.0752],
        [ 0.1106,  0.0752,  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.4116,  0.4618,  0.3337], requires_grad=True), 'parametrizations.weight.original': Parameter containing:
tensor([[ 0.0344,  0.0072,  0.0562],
        [ 0.2481,  0.2788,  0.4807],
        [ 0.3946, -0.2378,  0.0221]], requires_grad=True)}

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

print(layer.parametrizations.weight.original)

Parameter containing:
tensor([[ 0.0344,  0.0072,  0.0562],
        [ 0.2481,  0.2788,  0.4807],
        [ 0.3946, -0.2378,  0.0221]], 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。

快取引數化值#

引數化透過上下文管理器 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

連線引數化#

連線兩個引數化就像在同一個張量上註冊它們一樣簡單。我們可以使用它來從更簡單的引數化建立更復雜的引數化。例如，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(1.9881e-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(2.5723e-07, grad_fn=<DistBackward0>)
tensor(5.2684e-09, grad_fn=<DistBackward0>)
tensor(True)

初始化引數化#

引數化提供了一種初始化它們的方法。如果我們實現一個具有以下簽名的 right_inverse 方法：

def right_inverse(self, X: 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(0.1324, grad_fn=<DistBackward0>)

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

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

此方法的名稱來源於我們通常期望 forward(right_inverse(X)) == X。這是重寫使用值 X 初始化後的前向傳播應返回值 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 中的零位置與掩碼中的零位置相同時才成立。即使那樣，如果我們為一個剪枝引數分配一個張量，那麼該張量實際上已經被剪枝也就不足為奇了。

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.1694, 0.1887, 0.3677],
        [0.4180, 0.1883, 0.1400],
        [0.9703, 0.4129, 0.7185],
        [0.5919, 0.7431, 0.2885]])

Initialized weight:
tensor([[0.1694, 0.1887, 0.3677],
        [0.0000, 0.1883, 0.1400],
        [0.9703, 0.0000, 0.7185],
        [0.5919, 0.7431, 0.2885]], 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.5033, -0.5729,  0.0168],
        [ 0.0554,  0.1867,  0.5595],
        [-0.4279,  0.0124, -0.4335]], requires_grad=True)

Parametrized:
ParametrizedLinear(
  in_features=3, out_features=3, bias=True
  (parametrizations): ModuleDict(
    (weight): ParametrizationList(
      (0): Skew()
    )
  )
)
tensor([[ 0.0000, -0.5729,  0.0168],
        [ 0.5729,  0.0000,  0.5595],
        [-0.0168, -0.5595,  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.5729,  0.0168],
        [ 0.5729,  0.0000,  0.5595],
        [-0.0168, -0.5595,  0.0000]], requires_grad=True)

移除引數化時，我們可以選擇保留原始引數（即 layer.parametriations.weight.original 中的引數），而不是其引數化版本，方法是將標誌 leave_parametrized=False 設定為 True。

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.5575,  0.0889, -0.4280],
        [-0.1325, -0.5174,  0.2741],
        [-0.2705,  0.4845, -0.5173]], requires_grad=True)

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

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

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

在此頁面上
PyTorch 庫