注意

点击这里下载完整示例代码

参数化教程¶

创建日期：2021年4月19日 | 最后更新日期：2024年2月5日 | 最后验证日期：2024年11月5日

深度学习模型的正则化是一项 surprisingly 具有挑战性的任务。诸如惩罚方法等经典技术在应用于深度模型时，由于优化函数的复杂性，往往效果不佳。当处理病态模型时，这个问题尤为严重。这些模型的例子包括在长序列上训练的 RNN 和 GAN。近年来，已经提出了一些技术来正则化这些模型并提高它们的收敛性。对于递归模型，有人建议控制 RNN 的递归核的奇异值，以使其具有良好的条件数。例如，可以通过使递归核正交来实现这一点。另一种正则化递归模型的方法是通过 “权重归一化”。这种方法建议将参数的学习与其范数的学习解耦。为此，将参数除以其 Frobenius 范数，并学习一个单独的参数来编码其范数。类似地，GAN 也提出了类似的正则化方法，称为 “谱归一化”。该方法通过将其参数除以它们的谱范数，而不是 Frobenius 范数，来控制网络的 Lipschitz 常数。

所有这些方法都有一个共同的模式：它们在使用参数之前都会以适当的方式对其进行转换。在第一种情况下，它们通过一个将矩阵映射到正交矩阵的函数使其正交。在权重归一化和谱归一化的情况下，它们将原始参数除以其范数。

更一般地说，所有这些示例都使用一个函数来为参数添加额外的结构。换句话说，它们使用一个函数来约束参数。

在这个教程中，你将学习如何实现和使用此模式来对你的模型施加约束。这样做简单得就像编写你自己的 nn.Module。

Requirements: 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))

此实现虽然正确且完整，但存在一些问题：

它重新实现了该层。我们必须将线性层实现为x @ A。对于线性层来说这并不是很大问题，但想象一下需要重新实现一个CNN或Transformer……
它不将层和参数化分开。如果参数化更加复杂，我们就必须为每个要使用的层重写其代码。
它每次使用该层时都会重新计算参数化。如果我们在一个前向传递过程中多次使用该层（想象一下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>)

检查一个参数化的模块¶

当一个模块被参数化时，我们发现该模块发生了三种变化：

model.weight 现在是一个属性
它有一个新的 module.parametrizations 属性
未参数化的权重已移动到 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。

参数化值的缓存¶

参数化带有通过上下文管理器内置的缓存系统 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(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

它将在分配到参数化张量时使用。

让我们将 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 的前向传播应返回值 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.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.parametriations.weight.original）而不是其参数化版本，通过设置标志 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.053 秒)

通过 Sphinx-Gallery 生成的画廊