目录
  • 教程 >
  • 参数化教程
快捷键
intermediate/参数化

备注

点击此处下载完整示例代码

参数化教程 ¶

创建于：2025 年 4 月 1 日 | 最后更新：2025 年 4 月 1 日 | 最后验证：2024 年 11 月 5 日

作者：马里奥·莱赞科

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

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

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

在本教程中，您将学习如何实现和使用这种模式来对您的模型施加约束。这样做就像编写自己的 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

我们可以用同样的方法处理任何其他层。例如，我们可以创建一个具有斜对称核的 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 属性。这个属性会在我们请求 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()

参数化的连接

将两个参数化拼接起来就像在同一个张量上注册它们一样简单。我们可以利用这一点从简单的参数化创建更复杂的参数化。例如，凯莱映射将反对称矩阵映射到正行列式的正交矩阵。我们可以将 Skew 和实现凯莱映射的参数化拼接起来，得到一个具有正交权重的层

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 的表现符合预期。为了看到这一点，让我们升级凯莱参数化以支持初始化。

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 在相同的位置有零时，这个条件才成立。即使在这种情况下，如果我们把一个张量分配给剪枝参数，那么这个张量实际上被剪枝也就不足为奇了

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 秒）

由 Sphinx-Gallery 生成的画廊

评分这个教程

© 版权所有 2024，PyTorch。

使用 Sphinx 构建，主题由 Read the Docs 提供。
//暂时添加调查链接

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取初学者和高级开发者的深入教程

查看教程

资源

查找开发资源并获得您的疑问解答

查看资源

为了分析流量并优化您的体验，我们在本网站上提供 cookies。通过点击或导航，您同意允许我们使用 cookies。作为本站点的当前维护者，适用 Facebook 的 Cookies 政策。了解更多信息，包括可用的控制选项：Cookies 政策。