torch.nn.utils.convert_conv2d_weight_memory_format

torch.nn.utils.convert_conv2d_weight_memory_format(module, memory_format)

将 nn.Conv2d.weight 的 memory_format 转换为指定的 memory_format。

此转换递归应用于嵌套的 nn.Module，包括 module 本身。请注意，它仅更改 memory_format，而不改变每个维度的语义。此函数用于促进计算采用 NHWC 内核，这可以在计算能力 >= 7.0 的 CUDA 设备上为 fp16 数据带来显著的加速。

注意

调用 model.to(memory_format=torch.channels_last) 比工具函数 convert_conv2d_weight_memory_format 更激进。任何具有 4D 权重的层都会受到 model.to 的影响，而这些层不一定能从转换为指定的 memory_format 中受益。我们确信的一点是，在 cuDNN 中对卷积进行 NHWC (channels_last) 转换是有益的，因为它有利于在 NHWC 中运行卷积，即使在必须对输入张量应用置换 (permutation) 的情况下也是如此。

因此，我们的策略是只将卷积的权重转换为 channels_last。这确保了：1. 将使用快速卷积内核，其收益可能超过置换 (permutation) 的开销（如果输入格式不同）。2. 在不会从 memory_format 转换中受益的层上不会应用不必要的置换。

最佳情况是，卷积层之间的层与 channels last 兼容。输入张量在遇到第一个卷积层时会被置换为 channels last 格式并保持该内存格式。因此，后续的卷积无需对其输入张量进行置换。

如果卷积层之间存在与 channels last 不兼容的层，我们需要将输入张量对该层置换回 contiguous format。输入张量将以 contiguous format 通过剩余的层，并在遇到另一个卷积层时被置换为 channels last 格式。将该置换传播到更早的层没有意义，因为大多数层对 memory_format 相当不敏感。

当 PyTorch 支持置换融合时，这种说法可能会改变，因为可能存在比紧邻卷积之前更好的位置来融合置换。

参数

• module (nn.Module) – nn.Conv2d & nn.ConvTranspose2d 或容器 nn.Module

• memory_format – 用户指定的 memory_format，例如 torch.channels_last 或 torch.contiguous_format

返回

更新了 nn.Conv2d 的原始模块

示例

>>> input = torch.randint(1, 10, (2, 8, 4, 4), dtype=torch.float16, device="cuda")
>>> model = nn.Sequential(
>>>     nn.Conv2d(8, 4, 3)).cuda().half()
>>> # This is identical to:
>>> # nn.utils.convert_conv2d_weight_memory_format(model, torch.channels_last)
>>> model = nn.utils.convert_conv2d_weight_memory_format(model, torch.channels_last)
>>> out = model(input)