torch.nn.utils.convert_conv3d_weight_memory_format¶

torch.nn.utils.convert_conv3d_weight_memory_format(module, memory_format)[源代码]¶

将 nn.Conv3d.weight 的 memory_format 转换为 memory_format 该转换递归地应用于嵌套的 nn.Module，包括 module。请注意，它仅更改 memory_format，但不改变每个维度的语义。此函数用于促进计算以采用 NHWC 内核，这为计算能力 >= 7.0 的 CUDA 设备上的 fp16 数据提供了显著的速度提升。

注意

调用 model.to(memory_format=torch.channels_last) 比实用函数 convert_conv3d_weight_memory_format 更为激进。任何具有4D权重的层都会受到 model.to 的影响，这并不一定会从转换到指定的 memory_format 中受益。我们确信的一个地方是，在cuDNN中进行卷积的NHWC（channels_last）转换，因为即使在必须对输入张量应用排列的情况下，NHWC运行卷积也是有利的。

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

最佳情况是，卷积层之间的层是通道最后兼容的。输入张量在遇到第一个卷积层时会被置换为通道最后格式，并保持在该内存格式中。因此，后续的卷积操作不需要再置换其输入张量。

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

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

Parameters

模块 (nn.Module) – nn.Conv3d & nn.ConvTranspose3d 或容器 nn.Module
memory_format – 用户指定的 memory_format，例如 torch.channels_last 或 torch.contiguous_format

Returns

更新后的原始模块，包含 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()
>>> # 这等同于：
>>> # nn.utils.convert_conv3d_weight_memory_format(model, torch.channels_last)
>>> model = nn.utils.convert_conv3d_weight_memory_format(model, torch.channels_last)
>>> out = model(input)