参数化教程
原文：
译者：飞龙
注意
点击这里下载完整示例代码
在本教程中，您将学习如何实现并使用此模式来对模型进行约束。这样做就像编写自己的nn.Module一样容易。
对深度学习模型进行正则化是一项令人惊讶的挑战。传统技术，如惩罚方法，通常在应用于深度模型时效果不佳，因为被优化的函数的复杂性。当处理病态模型时，这一点尤为棘手。这些模型的示例包括在长序列上训练的 RNN 和 GAN。近年来已经提出了许多技术来对这些模型进行正则化并改善它们的收敛性。对于循环模型，已经提出控制 RNN 的循环核的奇异值以使其具有良好条件性。例如，可以通过使循环核正交来实现这一点。另一种正则化循环模型的方法是通过“权重归一化”。该方法建议将参数的学习与其范数的学习分离。为此，将参数除以其Frobenius 范数，并学习一个编码其范数的单独参数。类似的正则化方法也适用于以“谱归一化”命名的 GAN。该方法通过将网络的参数除以其谱范数而不是其 Frobenius 范数来控制网络的 Lipschitz 常数。
所有这些方法都有一个共同的模式：它们在使用参数之前以适当的方式转换参数。在第一种情况下，它们通过使用将矩阵映射到正交矩阵的函数使其正交。在权重和谱归一化的情况下，它们通过将原始参数除以其范数来实现。
更一般地，所有这些示例都使用一个函数在参数上添加额外的结构。换句话说，它们使用一个函数来约束参数。
它没有将层和参数化分开。如果参数化更加困难，我们将不得不为要在其中使用它的每个层重新编写其代码。
要求：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))
尽管这种实现是正确且独立的，但存在一些问题：
- 它重新实现了该层。我们必须将线性层实现为
x @ A。对于线性层来说，这并不是非常困难，但想象一下必须重新实现 CNN 或 Transformer… - 手动实现参数化
- 每次使用该层时都会重新计算参数化。如果在前向传递期间多次使用该层（想象一下 RNN 的循环核），它将在每次调用该层时计算相同的
A。
原文：pytorch.org/tutorials/intermediate/parametrizations.html
参数化可以解决所有这些问题以及其他问题。
让我们从使用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
我们可以对任何其他层执行相同的操作。例如，我们可以创建一个具有斜对称核的 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])
检查参数化模块
当一个模块被参数化时，我们发现模块以三种方式发生了变化：
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}")
这个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 映射将斜对称矩阵映射到正行列式的正交矩阵。我们可以连接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 # (X - I)(X + I)^{-1} return torch.linalg.solve(X + self.Id, self.Id - X) 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在与掩码相同的位置有零时才成立。即使是这样，如果我们将一个张量分配给一个被修剪的参数，那么这个张量实际上将被修剪也就不足为奇了
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)
脚本的总运行时间：（0 分钟 0.000 秒）
下载 Python 源代码：parametrizations.py
下载 Jupyter 笔记本：parametrizations.ipynb
PyTorch 2.2 中文官方教程（十四）（2）https://developer.aliyun.com/article/1482571
