目录
parametrize.register_parametrization
parametrize.remove_parametrizations
parametrize.ParametrizationList
torch.nn.utils.parametrize.register_parametrization是PyTorch中的一个功能，它允许用户将自定义参数化方法应用于模块中的张量。这种方法对于改变和控制模型参数的行为非常有用，特别是在需要对参数施加特定的约束或转换时。
nn.Module相关联，可以对其行为进行自定义。module.parametrizations.[tensor_name].original访问原始张量，并通过module.[tensor_name]访问参数化后的版本。cached()上下文管理器来激活，以提高效率。right_inverse方法，可以自定义参数化的初始值。module (nn.Module): 需要注册参数化的模块。tensor_name (str): 需要进行参数化的参数或缓冲区的名称。parametrization (nn.Module): 将要注册的参数化。unsafe (bool, 可选): 表示参数化是否可能改变张量的数据类型和形状。默认为False。unsafe=True，则在注册时不会检查参数化的一致性，这可能带来风险。tensor_name的参数或缓冲区，将抛出ValueError。import torch
import torch.nn as nn
import torch.nn.utils.parametrize as P
# 定义一个对称矩阵参数化
class Symmetric(nn.Module):
def forward(self, X):
return X.triu() + X.triu(1).T
def right_inverse(self, A):
return A.triu()
# 应用参数化
m = nn.Linear(5, 5)
P.register_parametrization(m, "weight", Symmetric())
print(torch.allclose(m.weight, m.weight.T)) # 现在m.weight是对称的
# 初始化对称权重
A = torch.rand(5, 5)
A = A + A.T
m.weight = A
print(torch.allclose(m.weight, A))
这个示例创建了一个线性层，对其权重应用了对称性参数化，然后初始化权重为一个对称矩阵。通过这种方法，可以确保模型的权重始终保持特定的结构特性。
torch.nn.utils.parametrize.remove_parametrizations 是 PyTorch 中的一个功能，它用于移除模块中某个张量上的参数化。这个函数允许用户将模块中的参数从参数化状态恢复到原始状态，根据leave_parametrized参数的设置，可以选择保留当前参数化的输出或恢复到未参数化的原始张量。
module (nn.Module): 从中移除参数化的模块。tensor_name (str): 要移除参数化的张量的名称。leave_parametrized (bool, 可选): 是否保留属性tensor_name作为参数化的状态。默认为True。module[tensor_name]未被参数化，会抛出ValueError。leave_parametrized=False且参数化依赖于多个张量，也会抛出ValueError。import torch
import torch.nn as nn
import torch.nn.utils.parametrize as P
# 定义模块和参数化
m = nn.Linear(5, 5)
P.register_parametrization(m, "weight", ...)
# 假设在这里进行了一些操作
# 移除参数化，保留当前参数化的输出
P.remove_parametrizations(m, "weight", leave_parametrized=True)
# 或者，移除参数化，恢复到原始未参数化的张量
P.remove_parametrizations(m, "weight", leave_parametrized=False)
?这个示例展示了如何在一个线性层上注册并最终移除参数化。根据leave_parametrized的设置，可以选择在移除参数化后保留当前的参数化状态或恢复到原始状态。这使得在模型开发和实验过程中可以更灵活地控制参数的行为。
torch.nn.utils.parametrize.cached() 是 PyTorch 框架中的一个上下文管理器，用于启用通过 register_parametrization() 注册的参数化对象的缓存系统。当这个上下文管理器活跃时，参数化对象的值在第一次被请求时会被计算和缓存。离开上下文管理器时，缓存的值会被丢弃。
P.cached() 的上下文管理器内来激活缓存。import torch.nn as nn
import torch.nn.utils.parametrize as P
class MyModel(nn.Module):
# 模型定义
...
model = MyModel()
# 应用一些参数化
...
# 使用缓存系统包装模型的前向传播
with P.cached():
output = model(inputs)
# 或者，仅在特定部分使用缓存
with P.cached():
for x in xs:
out_rnn = self.rnn_cell(x, out_rnn)
?这个示例展示了如何在模型的整个前向传播过程中或者在特定部分（如RNN循环中）使用缓存系统。这样做可以在保持模型逻辑不变的同时，提高计算效率。特别是在复杂的参数化场景中，这可以显著减少不必要的重复计算。
torch.nn.utils.parametrize.is_parametrized 是 PyTorch 库中的一个函数，用于检查一个模块是否有活跃的参数化，或者指定的张量名称是否已经被参数化。
module (nn.Module): 要查询的模块。tensor_name (str, 可选): 模块中要查询的属性，默认为None。如果提供，函数将检查此特定属性是否已经被参数化。import torch.nn as nn
import torch.nn.utils.parametrize as P
class MyModel(nn.Module):
# 模型定义
...
model = MyModel()
# 对模型的某个属性应用参数化
P.register_parametrization(model, 'weight', ...)
# 检查整个模型是否被参数化
is_parametrized = P.is_parametrized(model)
print(is_parametrized) # 输出 True 或 False
# 检查模型的特定属性是否被参数化
is_weight_parametrized = P.is_parametrized(model, 'weight')
print(is_weight_parametrized) # 输出 True 或 False
在这个示例中，is_parametrized 函数用来检查整个模型是否有任何参数化，以及模型的weight属性是否被特定地参数化。这对于验证参数化是否正确应用或在调试过程中理解模型的当前状态非常有用。
ParametrizationList 是 PyTorch 中的一个类，它是一个顺序容器，用于保存和管理经过参数化的 torch.nn.Module 的原始参数或缓冲区。当使用 register_parametrization() 对模块中的张量进行参数化时，这个容器将作为 module.parametrizations[tensor_name] 的类型存在。
ParametrizationList 保存了原始的参数或缓冲区，这些参数或缓冲区通过参数化被修改。right_inverse 方法，这些张量将以 original0, original1, … 等的形式被保存。modules (sequence): 代表参数化的模块序列。original (Parameter or Tensor): 被参数化的参数或缓冲区。unsafe (bool): 表明参数化是否可能改变张量的数据类型和形状。默认为False。当unsafe=True时，不会在注册时检查参数化的一致性，使用时需要小心。right_inverse(value): 按照注册的相反顺序调用参数化的 right_inverse 方法。然后，如果 right_inverse 输出一个张量，就将结果存储在 self.original 中；如果输出多个张量，就存储在 self.original0, self.original1, … 中。register_parametrization() 内部使用，并不建议用户直接实例化。unsafe 参数的使用需要谨慎，因为它可能带来一致性问题。由于 ParametrizationList 主要用于内部实现，因此一般不会直接在用户代码中创建实例。它在进行参数化操作时自动形成，例如：
import torch.nn as nn
import torch.nn.utils.parametrize as P
# 定义一个简单的模型
class MyModel(nn.Module):
def __init__(self):
super(MyModel, self).__init__()
self.linear = nn.Linear(10, 10)
model = MyModel()
# 对模型的某个参数应用参数化
P.register_parametrization(model.linear, "weight", MyParametrization())
# ParametrizationList 实例可以通过以下方式访问
param_list = model.linear.parametrizations.weight
?在这个示例中，param_list 将是 ParametrizationList 类的一个实例，包含了 weight 参数的所有参数化信息。
本篇博客探讨了 PyTorch 中 torch.nn.utils.parametrize 子模块的强大功能和灵活性。它详细介绍了如何通过自定义参数化（register_parametrization）来改变和控制模型参数的行为，提供了移除参数化（remove_parametrizations）的方法以恢复模型到原始状态，并探讨了如何利用缓存机制（cached）来提高参数化参数在前向传播中的计算效率。此外，文章还解释了如何检查模型或其属性的参数化状态（is_parametrized），并深入了解了 ParametrizationList 类在内部如何管理参数化参数。