评价此页
★ ★ ★ ★ ★
intermediate/parametrizations

注意

跳转至页面底部 下载完整示例代码。

参数化教程#

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

作者: Mario Lezcano

深度学习模型的正则化是一项极具挑战性的任务。传统的惩罚方法在应用于深度模型时，由于优化函数的复杂性，往往效果不佳。这在处理病态模型（ill-conditioned models）时尤为突出，例如在长序列上训练的 RNN 以及 GAN。近年来，人们提出了多种技术来正则化这些模型并改善其收敛性。在循环模型上，有人提出控制循环核（recurrent kernel）的奇异值，以使 RNN 达到良好的条件数。例如，可以通过使循环核正交化来实现。另一种正则化循环模型的方法是“权重归一化”（weight normalization）。该方法提议将参数的学习与其范数的学习分离开来。为此，参数除以其Frobenius 范数，并学习一个独立的参数来编码其范数。针对 GAN，有人提出了类似的正则化方法，称为“谱归一化”（spectral normalization）。该方法通过将参数除以其谱范数而非 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.8036, 0.3987, 0.1443],
        [0.3987, 0.0383, 0.6465],
        [0.1443, 0.6465, 0.5769]])

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

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.3249,  0.1577,  0.0665],
        [ 0.1577,  0.1411, -0.0831],
        [ 0.0665, -0.0831,  0.3240]], 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.1352, -0.0479],
        [-0.1352,  0.0000,  0.1349],
        [ 0.0479, -0.1349,  0.0000]], grad_fn=<SelectBackward0>)
tensor([[ 0.0000, -0.1031,  0.0231],
        [ 0.1031,  0.0000, -0.0859],
        [-0.0231,  0.0859,  0.0000]], grad_fn=<SelectBackward0>)

检查已参数化的模块#

当一个模块被参数化后，我们会发现该模块在三个方面发生了变化

  1. model.weight 现在是一个属性（property）

  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.3505, 0.1270, 0.1461], requires_grad=True), 'parametrizations.weight.original': Parameter containing:
tensor([[-0.5666,  0.2890,  0.2447],
        [ 0.4597, -0.0548,  0.3801],
        [ 0.4598,  0.1729, -0.4816]], requires_grad=True)}

它现在位于 layer.parametrizations.weight.original 下

print(layer.parametrizations.weight.original)

Parameter containing:
tensor([[-0.5666,  0.2890,  0.2447],
        [ 0.4597, -0.0548,  0.3801],
        [ 0.4598,  0.1729, -0.4816]], 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 map）将斜对称矩阵映射为行列式为正的正交矩阵。我们可以级联 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

tensor(2.3935e-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.1900e-07, grad_fn=<DistBackward0>)
tensor(2.6342e-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 也能按预期工作。为了验证这一点，让我们升级凯莱参数化，使其也支持初始化

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(1.9016, 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.9971, 0.1024, 0.1147],
        [0.5218, 0.3077, 0.2431],
        [0.1137, 0.1380, 0.9402],
        [0.5013, 0.5014, 0.6396]])

Initialized weight:
tensor([[0.9971, 0.0000, 0.0000],
        [0.5218, 0.3077, 0.2431],
        [0.1137, 0.1380, 0.9402],
        [0.5013, 0.5014, 0.6396]], 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.3509, -0.0763,  0.0548],
        [-0.2591,  0.1056, -0.2055],
        [ 0.2804,  0.2821, -0.2839]], requires_grad=True)

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

在移除参数化时，我们可以通过设置标志 leave_parametrized=False，选择保留原始参数（即 layer.parametriations.weight.original 中的参数），而不是其参数化后的版本

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.2105, -0.3832, -0.5662],
        [-0.5034, -0.0268, -0.2557],
        [ 0.1657, -0.5457, -0.5658]], requires_grad=True)

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

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

脚本运行总耗时： (0 分 0.054 秒)

在此页面上
PyTorch 库