参数化教程

Created On: Apr 19, 2021 | Last Updated: Feb 05, 2024 | Last Verified: Nov 05, 2024

作者： Mario Lezcano

对深度学习模型进行正则化是一项非常具有挑战性的任务。当应用于深度模型时，由于优化函数的复杂性，经典的正则化技术（例如惩罚方法）通常表现不佳。这在处理病态模型时尤其成问题。这类模型的例子包括在长序列上训练的 RNN 和 GAN。近年来提出了许多技术来正则化这些模型并改善其收敛性。对于循环模型，已经建议控制循环核的奇异值以使 RNN 成为病态良好的模型。例如，这可以通过使循环核为`正交矩阵 <https://en.wikipedia.org/wiki/Orthogonal_matrix>`_ 来实现。对循环模型进行正则化的另一种方法是 “权重归一化”。该方法建议将参数的学习与其模的学习进行解耦。为此，将参数除以其 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

然后我们可以使用这个想法来实现一个具有对称权重的线性层。

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())

现在，线性层的矩阵是对称的。

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

我们可以对任何其他层做同样的事情。例如，我们可以创建一个具有`反对称内核 <https://en.wikipedia.org/wiki/Skew-symmetric_matrix>`_的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])

检验参数化的模块

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

1. ``model.weight``现在是一个属性。

2. 它有了一个新的``module.parametrizations``属性。

3. 未参数化的权重已被移动到``module.parametrizations.weight.original``。

在参数化``weight``之后，layer.weight``变成了一个`Python属性 <https://docs.python.org/3/library/functions.html#property>`_。此属性每次请求``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}")

此``parametrizations``属性是一个``nn.ModuleDict``，可以像这样访问它。

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

此``nn.ModuleDict``的每个元素都是一个``ParametrizationList``，它的行为类似于``nn.Sequential``。此列表允许我们在一个权重上级联参数化。由于这是一个列表，我们可以通过索引访问参数化。这是我们``Symmetric``参数化所在的位置。

print(layer.parametrizations.weight[0])

我们注意到的另一件事是，如果我们打印参数，我们会看到参数``weight``已被移到：

print(dict(layer.named_parameters()))

现在位于``layer.parametrizations.weight.original``下。

print(layer.parametrizations.weight.original)

除了这三个小变化，参数化与我们的手动实现完全相同。

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

参数化是一级公民

由于``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()

串联参数化

串联两个参数化就像在同一个张量上注册它们一样简单。我们可以使用它从更简单的参数化创建更复杂的参数化。例如，`Cayley映射 <https://en.wikipedia.org/wiki/Cayley_transform#Matrix_map>`_将反对称矩阵映射到正定确定性的正交矩阵。我们可以串联``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

这也可以用来修剪参数化的模块或重用参数化。例如，矩阵指数将对称矩阵映射到对称正定(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

初始化参数化

参数化附带一个初始化机制。如果我们实现了一个具有以下签名的方法``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

当我们串联参数化时，这个``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

这个初始化步骤可以更简洁地写为：

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``在掩码(Mask)相同位置具有零值时才为真。即便如此，如果我们将一个张量分配给一个修剪参数，那么毫不奇怪，这个张量实际上将被修剪。

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}")

移除参数化

我们可以通过使用``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)

在移除参数化时，我们可以选择是否保留原始参数（即位于``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)