torch.nn.utils.convert_conv3d_weight_memory_format

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

将 memory_format 的 nn.Conv3d.weight 转换为 memory_format 。转换会递归地应用于嵌套的 nn.Module ，包括 module 。请注意，它只更改 memory_format，但不更改每个维度的语义。此函数用于简化计算以采用 NHWC 内核，这为 CUDA 设备上的 fp16 数据提供了相当大的加速，设备计算能力需>= 7.0

注意

调用 model.to(memory_format=torch.channels_last_3d) 比实用函数 convert_conv3d_weight_memory_format 更具侵略性。任何具有 4d 权重的层都会受到 model.to 的影响，这并不一定从转换为指定的 memory_format 中受益。我们非常有信心的一点是，在 cuDNN 中进行卷积的 NDHWC(channels_last_3d) 转换，因为即使在需要应用排列输入张量的情况下，运行卷积在 NDHWC 中也是有利的。

因此，我们的策略是仅将卷积的权重转换为 channels_last_3d。这确保了：1. 将使用快速卷积内核，其好处可能超过排列的开销（如果输入不是相同的格式）。2. 不对那些从内存格式转换中受益的层应用不必要的排列。

最佳情况是，卷积层之间的层是 channels last 兼容的。当输入张量遇到第一个卷积层时，它将被排列为 channels last 格式，并保持在该内存格式。因此，后续的卷积不需要对其输入张量进行排列。

在通道的最后一个不兼容层位于卷积层之间的情况下，我们需要将该层的输入张量重新排列回连续格式。输入张量将以连续格式通过剩余的层，并在遇到另一个卷积层时再次重新排列为通道最后格式。将这种排列传播到更早的层没有意义，因为大多数层对 memory_format 相当不敏感。

当 PyTorch 支持排列融合时，这种说法可能会改变，因为可能存在比立即在卷积层之前融合排列更好的位置。

参数:

• 模块（nn.Module）- nn.Conv3d & nn.ConvTranspose3d 或容器 nn.Module

• 内存格式 - 用户指定 memory_format ，例如 torch.channels_last 或 torch.contiguous_format

返回值:

带有更新 nn.Conv3d 的原始模块

示例

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