参数化教程

创建时间： 2021年4月19日 |上次更新时间：2024 年 2 月 5 日 |上次验证： Nov 05， 2024

作者： Mario Lezcano

正则化深度学习模型是一项极具挑战性的任务。 惩罚法等经典技术在应用时往往达不到要求 在深度模型上，因为正在优化的函数的复杂性。 这在处理病态模型时尤其成问题。 这些示例是在长序列上训练的 RNN 和 GAN。一个数字 近年来提出了一些技术来规范这些 模型并提高它们的收敛性。在循环模式中，情况一直如此 建议控制 RNN 的条件良好。这可以通过实现，例如，使 循环核正交。 正则化递归模型的另一种方法是通过 “权重标准化”。 这种方法建议将参数的学习与 学习他们的规范。为此，将参数除以 Frobenius 范数，并学习一个单独的参数来编码其范数。 为 GAN 提出了类似的正则化，名称为 “光谱归一化”。此方法 通过将网络的参数除以 它们的光谱范数， 而不是他们的 Frobenius 规范。

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

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

在本教程中，您将学习如何实现和使用此模式将 constraints 的这样做就像编写自己的 .nn.Module

要求：torch>=1.9.0

手动实现参数化

假设我们想要一个权重对称的方形线性层，即 权重使得 。一种方法是 将矩阵的上三角部分复制到其下三角部分XX = 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))

这种实现虽然正确且自包含，但也存在许多问题：

1. 它将重新实现该层。我们必须将线性层实现为 .这是 对于线性层来说不是很有问题，但想象一下必须重新实现 CNN 或 变压器。。。x @ A

2. 它不会将层和参数化分开。如果参数化为 更困难的是，我们将不得不为我们想要使用它的每个层重写它的代码 在。

3. 每次我们使用 layer 时，它都会重新计算参数化。如果我们使用 layer 在前向传递期间多次（想象一下 RNN 的递归内核），它 每次调用该层时，都会计算相同的结果。A

参数化简介

参数化可以解决所有这些问题以及其他问题。

让我们首先使用 . 我们唯一需要做的就是将参数化编写为常规torch.nn.utils.parametrizenn.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>)

检查参数化模块

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

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

2. 它有一个新属性module.parametrizations

3. 未参数化的权重已移至module.parametrizations.weight.original

参数化后，将转换为 Python 属性。 每次我们请求时，此属性都会计算，就像我们在上述实现中所做的那样。weightlayer.weightparametrization(weight)layer.weightLinearSymmetric

已注册的参数化存储在模块内的 attribute 下。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()
    )
  )
)

此属性是一个 ，可以按此方式访问parametrizationsnn.ModuleDict

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

ModuleDict(
  (weight): ParametrizationList(
    (0): Symmetric()
  )
)
ParametrizationList(
  (0): Symmetric()
)

this 的每个元素都是一个 ，其行为类似于 .此列表将允许我们在一个权重上连接参数化。 由于这是一个列表，我们可以访问为其编制索引的参数化。这是 我们的参数化所在nn.ModuleDictParametrizationListnn.SequentialSymmetric

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

参数化是一等公民

由于 是 ，这意味着参数化 被正确注册为原始模块的子模块。因此，相同的规则 for register parameters in a module apply 来注册参数化。 例如，如果参数化具有参数，则这些参数将从 CPU 中移出 到 CUDA 中调用 .layer.parametrizationsnn.ModuleListmodel = model.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 映射将偏对称矩阵映射到正行列式矩阵的正交矩阵。我们可以 concatenate 和一个参数化，该参数化实现 Cayley 映射以获取具有 正交权重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.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

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

让我们升级 class 的实现来支持这一点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>)

当我们连接参数化时，这按预期工作。 为了实现这一点，让我们升级 Cayley 参数化以支持初始化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(2.2141, grad_fn=<DistBackward0>)

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

layer_orthogonal.weight = nn.init.orthogonal_(layer_orthogonal.weight)

这种方法的名称来源于我们经常期望的事实 那。这是重写 初始化后的 forward with 值应返回值 . 这种约束在实践中并没有得到强有力的执行。事实上，有时，它可能是 interest 来放松这种关系。例如，请考虑以下实现 随机修剪方法：forward(right_inverse(X)) == XXX

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 . 仅当矩阵在与掩码相同的位置有零时，这才成立。 即便如此，如果我们将张量分配给修剪后的参数，也就不足为奇了 实际上，该 Tensor 将被修剪forward(right_inverse(A)) == AA

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.originalleave_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)