torch.nn.utils.convert_conv3d_weight_memory_format¶
- torch.nn.utils.convert_conv3d_weight_memory_format(module, memory_format)[source]¶
将
nn.Conv3d.weight
的memory_format
转换为memory_format
此转换递归地应用于嵌套的nn.Module
,包括module
。请注意,它只更改 memory_format,而不更改每个维度的语义。此函数用于促进计算以采用 NHWC 内核,这为具有计算能力 >= 7.0 的 CUDA 设备上的 fp16 数据提供了相当大的加速。注意
调用
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. 对那些不会从 memory_format 转换中受益的层不会进行不必要的置换。
最佳情况是,卷积层之间的层与 channels last 兼容。当遇到第一个卷积层时,输入张量将被置换为 channels last,并保持该内存格式。因此,后续卷积将不需要置换其输入张量。
如果卷积层之间存在与 channels last 不兼容的层,我们需要将输入张量置换回连续格式以用于该层。输入张量将以连续格式通过剩余的层,并在遇到另一个卷积层时被置换为 channels last。将该置换传播到更早的层没有意义,因为大多数层对
memory_format
非常不敏感。当 PyTorch 支持置换融合时,这种说法可能会改变,因为可能在卷积之前的某个更好位置融合置换。
- 参数
module (nn.Module) –
nn.Conv3d
和nn.ConvTranspose3d
或容器nn.Module
memory_format – 用户指定的
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)