目录
  • 教程 >
  • (beta) PyTorch 中的通道最后内存格式
快捷键
中间件/内存格式教程

备注

点击此处下载完整示例代码

(beta) PyTorch 中通道最后内存格式 ¶

创建于：2025 年 4 月 1 日 | 最后更新：2025 年 4 月 1 日 | 最后验证：2024 年 11 月 5 日

作者：Vitaly Fedyunin

什么是通道最后 ¶

通道最后内存格式是一种在内存中保持维度顺序的 NCHW 张量的替代排序方式。通道最后张量的排序方式使得通道成为最密集的维度（即按像素逐像素存储图像）。

例如，经典（连续）存储的 NCHW 张量（在我们的例子中是两个 4x4 图像，具有 3 个颜色通道）看起来是这样的：

classic_memory_format

通道最后内存格式对数据的排序方式不同：

channels_last_memory_format

PyTorch 通过利用现有的步长结构支持内存格式（并提供对现有模型（包括 eager、JIT 和 TorchScript）的向后兼容性）。例如，10x3x16x16 批次的通道最后格式将具有步长等于（768，1，48，3）。

仅实现了 4D NCHW 张量的最后内存格式。

内存格式 API ¶

这里是如何在连续和通道最后内存格式之间转换张量的。

经典 PyTorch 连续张量

import torch

N, C, H, W = 10, 3, 32, 32
x = torch.empty(N, C, H, W)
print(x.stride())  # Outputs: (3072, 1024, 32, 1)

转换运算符

x = x.to(memory_format=torch.channels_last)
print(x.shape)  # Outputs: (10, 3, 32, 32) as dimensions order preserved
print(x.stride())  # Outputs: (3072, 1, 96, 3)

返回连续

x = x.to(memory_format=torch.contiguous_format)
print(x.stride())  # Outputs: (3072, 1024, 32, 1)

替代选项

x = x.contiguous(memory_format=torch.channels_last)
print(x.stride())  # Outputs: (3072, 1, 96, 3)

格式检查

print(x.is_contiguous(memory_format=torch.channels_last))  # Outputs: True

两个 API to 和 contiguous 之间有细微差别。我们建议在显式转换张量内存格式时坚持使用 to 。

对于一般情况，两个 API 的行为相同。然而，在特殊情况下，对于大小为 NCHW 的 4 维张量，当以下条件之一成立时： C==1 或 H==1 && W==1 ，只有 to 才能生成表示通道最后内存格式的正确步长。

这是因为在上述两种情况中，张量的内存格式是模糊的，即大小为 N1HW 的连续张量既是 contiguous 也是内存存储中的通道最后。因此，它们已经被认为是给定内存格式的 is_contiguous ，因此 contiguous 调用成为空操作，不会更新步长。相反， to 会重新设置尺寸为 1 的维度的步长，以正确表示预期的内存格式。

special_x = torch.empty(4, 1, 4, 4)
print(special_x.is_contiguous(memory_format=torch.channels_last))  # Outputs: True
print(special_x.is_contiguous(memory_format=torch.contiguous_format))  # Outputs: True

同样，对于显式排列 API permute 也适用。在可能发生歧义的特殊情况下， permute 不能保证产生正确携带预期内存格式的步长。我们建议使用带有显式内存格式的 to 以避免意外行为。

在极端情况下，即三个非批量维度都等于 1 （ C==1 && H==1 && W==1 ），当前实现无法将张量标记为通道最后内存格式。

创建为通道最后

x = torch.empty(N, C, H, W, memory_format=torch.channels_last)
print(x.stride())  # Outputs: (3072, 1, 96, 3)

clone 保留内存格式

y = x.clone()
print(y.stride())  # Outputs: (3072, 1, 96, 3)

to ， cuda ， float … 保留内存格式

if torch.cuda.is_available():
    y = x.cuda()
    print(y.stride())  # Outputs: (3072, 1, 96, 3)

empty_like , *_like 运算符保留内存格式

y = torch.empty_like(x)
print(y.stride())  # Outputs: (3072, 1, 96, 3)

矩阵运算符保留内存格式

z = x + y
print(z.stride())  # Outputs: (3072, 1, 96, 3)

Conv , Batchnorm 模块使用 cudnn 后端支持通道最后（仅适用于 cuDNN >= 7.6）。卷积模块与二进制 p-wise 运算符不同，以通道最后作为主导内存格式。如果所有输入都在连续内存格式中，则运算符将输出为连续内存格式。否则，输出将采用通道最后内存格式。

if torch.backends.cudnn.is_available() and torch.backends.cudnn.version() >= 7603:
    model = torch.nn.Conv2d(8, 4, 3).cuda().half()
    model = model.to(memory_format=torch.channels_last)  # Module parameters need to be channels last

    input = torch.randint(1, 10, (2, 8, 4, 4), dtype=torch.float32, requires_grad=True)
    input = input.to(device="cuda", memory_format=torch.channels_last, dtype=torch.float16)

    out = model(input)
    print(out.is_contiguous(memory_format=torch.channels_last))  # Outputs: True

当输入张量达到不支持通道最后的运算符时，内核应自动应用排列以恢复输入张量的连续性。这引入了开销并阻止了通道最后内存格式的传播。尽管如此，它保证了正确的输出。

性能提升

通道最后内存格式优化在 GPU 和 CPU 上均可用。在 GPU 上，使用具有 Tensor Core 支持的 NVIDIA 硬件以降低精度运行时，性能提升最为显著。我们能够在使用‘AMP（自动混合精度）’训练脚本的情况下，与连续格式相比，实现超过 22%的性能提升。我们的脚本使用的是 NVIDIA 提供的 AMP https://github.com/NVIDIA/apex。

python main_amp.py -a resnet50 --b 200 --workers 16 --opt-level O2  ./data

# opt_level = O2
# keep_batchnorm_fp32 = None <class 'NoneType'>
# loss_scale = None <class 'NoneType'>
# CUDNN VERSION: 7603
# => creating model 'resnet50'
# Selected optimization level O2:  FP16 training with FP32 batchnorm and FP32 master weights.
# Defaults for this optimization level are:
# enabled                : True
# opt_level              : O2
# cast_model_type        : torch.float16
# patch_torch_functions  : False
# keep_batchnorm_fp32    : True
# master_weights         : True
# loss_scale             : dynamic
# Processing user overrides (additional kwargs that are not None)...
# After processing overrides, optimization options are:
# enabled                : True
# opt_level              : O2
# cast_model_type        : torch.float16
# patch_torch_functions  : False
# keep_batchnorm_fp32    : True
# master_weights         : True
# loss_scale             : dynamic
# Epoch: [0][10/125] Time 0.866 (0.866) Speed 230.949 (230.949) Loss 0.6735125184 (0.6735) Prec@1 61.000 (61.000) Prec@5 100.000 (100.000)
# Epoch: [0][20/125] Time 0.259 (0.562) Speed 773.481 (355.693) Loss 0.6968704462 (0.6852) Prec@1 55.000 (58.000) Prec@5 100.000 (100.000)
# Epoch: [0][30/125] Time 0.258 (0.461) Speed 775.089 (433.965) Loss 0.7877287269 (0.7194) Prec@1 51.500 (55.833) Prec@5 100.000 (100.000)
# Epoch: [0][40/125] Time 0.259 (0.410) Speed 771.710 (487.281) Loss 0.8285319805 (0.7467) Prec@1 48.500 (54.000) Prec@5 100.000 (100.000)
# Epoch: [0][50/125] Time 0.260 (0.380) Speed 770.090 (525.908) Loss 0.7370464802 (0.7447) Prec@1 56.500 (54.500) Prec@5 100.000 (100.000)
# Epoch: [0][60/125] Time 0.258 (0.360) Speed 775.623 (555.728) Loss 0.7592862844 (0.7472) Prec@1 51.000 (53.917) Prec@5 100.000 (100.000)
# Epoch: [0][70/125] Time 0.258 (0.345) Speed 774.746 (579.115) Loss 1.9698858261 (0.9218) Prec@1 49.500 (53.286) Prec@5 100.000 (100.000)
# Epoch: [0][80/125] Time 0.260 (0.335) Speed 770.324 (597.659) Loss 2.2505953312 (1.0879) Prec@1 50.500 (52.938) Prec@5 100.000 (100.000)

通过 --channels-last true 允许以通道最后格式运行模型，观察到 22%的性能提升。

python main_amp.py -a resnet50 --b 200 --workers 16 --opt-level O2 --channels-last true ./data

# opt_level = O2
# keep_batchnorm_fp32 = None <class 'NoneType'>
# loss_scale = None <class 'NoneType'>
#
# CUDNN VERSION: 7603
#
# => creating model 'resnet50'
# Selected optimization level O2:  FP16 training with FP32 batchnorm and FP32 master weights.
#
# Defaults for this optimization level are:
# enabled                : True
# opt_level              : O2
# cast_model_type        : torch.float16
# patch_torch_functions  : False
# keep_batchnorm_fp32    : True
# master_weights         : True
# loss_scale             : dynamic
# Processing user overrides (additional kwargs that are not None)...
# After processing overrides, optimization options are:
# enabled                : True
# opt_level              : O2
# cast_model_type        : torch.float16
# patch_torch_functions  : False
# keep_batchnorm_fp32    : True
# master_weights         : True
# loss_scale             : dynamic
#
# Epoch: [0][10/125] Time 0.767 (0.767) Speed 260.785 (260.785) Loss 0.7579724789 (0.7580) Prec@1 53.500 (53.500) Prec@5 100.000 (100.000)
# Epoch: [0][20/125] Time 0.198 (0.482) Speed 1012.135 (414.716) Loss 0.7007197738 (0.7293) Prec@1 49.000 (51.250) Prec@5 100.000 (100.000)
# Epoch: [0][30/125] Time 0.198 (0.387) Speed 1010.977 (516.198) Loss 0.7113101482 (0.7233) Prec@1 55.500 (52.667) Prec@5 100.000 (100.000)
# Epoch: [0][40/125] Time 0.197 (0.340) Speed 1013.023 (588.333) Loss 0.8943189979 (0.7661) Prec@1 54.000 (53.000) Prec@5 100.000 (100.000)
# Epoch: [0][50/125] Time 0.198 (0.312) Speed 1010.541 (641.977) Loss 1.7113249302 (0.9551) Prec@1 51.000 (52.600) Prec@5 100.000 (100.000)
# Epoch: [0][60/125] Time 0.198 (0.293) Speed 1011.163 (683.574) Loss 5.8537774086 (1.7716) Prec@1 50.500 (52.250) Prec@5 100.000 (100.000)
# Epoch: [0][70/125] Time 0.198 (0.279) Speed 1011.453 (716.767) Loss 5.7595844269 (2.3413) Prec@1 46.500 (51.429) Prec@5 100.000 (100.000)
# Epoch: [0][80/125] Time 0.198 (0.269) Speed 1011.827 (743.883) Loss 2.8196096420 (2.4011) Prec@1 47.500 (50.938) Prec@5 100.000 (100.000)

以下模型列表完全支持通道最后格式，并在 Volta 设备上显示出 8%-35%的性能提升： alexnet ， mnasnet0_5 ， mnasnet0_75 ， mnasnet1_0 ， mnasnet1_3 ， mobilenet_v2 ， resnet101 ， resnet152 ， resnet18 ， resnet34 ， resnet50 ， resnext50_32x4d ， shufflenet_v2_x0_5 ， shufflenet_v2_x1_0 ， shufflenet_v2_x1_5 ， shufflenet_v2_x2_0 ， squeezenet1_0 ， squeezenet1_1 ， vgg11 ， vgg11_bn ， vgg13 ， vgg13_bn ， vgg16 ， vgg16_bn ， vgg19 ， vgg19_bn ， wide_resnet101_2 ， wide_resnet50_2

以下列表中的模型完全支持 Channels last，并在 Intel(R) Xeon(R) Ice Lake（或更高版本）CPU 上显示出 26%-76%的性能提升： alexnet ， densenet121 ， densenet161 ， densenet169 ， googlenet ， inception_v3 ， mnasnet0_5 ， mnasnet1_0 ， resnet101 ， resnet152 ， resnet18 ， resnet34 ， resnet50 ， resnext101_32x8d ， resnext50_32x4d ， shufflenet_v2_x0_5 ， shufflenet_v2_x1_0 ， squeezenet1_0 ， squeezenet1_1 ， vgg11 ， vgg11_bn ， vgg13 ， vgg13_bn ， vgg16 ， vgg16_bn ， vgg19 ， vgg19_bn ， wide_resnet101_2 ， wide_resnet50_2

转换现有模型

Channels last 支持不受现有模型限制，任何模型都可以转换为 channels last，并在输入（或某些权重）格式正确后立即通过图传播格式。然而，并非所有操作符都完全转换为支持 channels last（通常返回连续输出）。在上面的示例中，不支持 channels last 的层将停止内存格式传播。尽管如此，由于我们已经将模型转换为 channels last 格式，这意味着每个卷积层，其 4 维权重在 channels last 内存格式中，将恢复 channels last 内存格式并从更快的内核中受益。

# Need to be done once, after model initialization (or load)
model = model.to(memory_format=torch.channels_last)  # Replace with your model

# Need to be done for every input
input = input.to(memory_format=torch.channels_last)  # Replace with your input
output = model(input)

然而，并非所有操作符都完全转换为支持 channels last（通常返回连续输出）。在上面的示例中，不支持 channels last 的层将停止内存格式传播。尽管如此，由于我们已经将模型转换为 channels last 格式，这意味着每个卷积层，其 4 维权重在 channels last 内存格式中，将恢复 channels last 内存格式并从更快的内核中受益。

但是不支持通道最后（channels last）的操作符确实会通过排列引入开销。如果您想提高转换后模型的性能，可以选择调查并识别模型中不支持通道最后（channels last）的操作符。

这意味着您需要将使用的操作符列表与支持的操作符列表进行验证 https://github.com/pytorch/pytorch/wiki/Operators-with-Channels-Last-support，或者将内存格式检查引入到动态执行模式中并运行您的模型。

执行以下代码后，如果操作符的输出与输入的内存格式不匹配，操作符将引发异常。

def contains_cl(args):
    for t in args:
        if isinstance(t, torch.Tensor):
            if t.is_contiguous(memory_format=torch.channels_last) and not t.is_contiguous():
                return True
        elif isinstance(t, list) or isinstance(t, tuple):
            if contains_cl(list(t)):
                return True
    return False


def print_inputs(args, indent=""):
    for t in args:
        if isinstance(t, torch.Tensor):
            print(indent, t.stride(), t.shape, t.device, t.dtype)
        elif isinstance(t, list) or isinstance(t, tuple):
            print(indent, type(t))
            print_inputs(list(t), indent=indent + "    ")
        else:
            print(indent, t)


def check_wrapper(fn):
    name = fn.__name__

    def check_cl(*args, **kwargs):
        was_cl = contains_cl(args)
        try:
            result = fn(*args, **kwargs)
        except Exception as e:
            print("`{}` inputs are:".format(name))
            print_inputs(args)
            print("-------------------")
            raise e
        failed = False
        if was_cl:
            if isinstance(result, torch.Tensor):
                if result.dim() == 4 and not result.is_contiguous(memory_format=torch.channels_last):
                    print(
                        "`{}` got channels_last input, but output is not channels_last:".format(name),
                        result.shape,
                        result.stride(),
                        result.device,
                        result.dtype,
                    )
                    failed = True
        if failed and True:
            print("`{}` inputs are:".format(name))
            print_inputs(args)
            raise Exception("Operator `{}` lost channels_last property".format(name))
        return result

    return check_cl


old_attrs = dict()


def attribute(m):
    old_attrs[m] = dict()
    for i in dir(m):
        e = getattr(m, i)
        exclude_functions = ["is_cuda", "has_names", "numel", "stride", "Tensor", "is_contiguous", "__class__"]
        if i not in exclude_functions and not i.startswith("_") and "__call__" in dir(e):
            try:
                old_attrs[m][i] = e
                setattr(m, i, check_wrapper(e))
            except Exception as e:
                print(i)
                print(e)


attribute(torch.Tensor)
attribute(torch.nn.functional)
attribute(torch)

如果您发现不支持通道最后（channels last）张量的操作符并希望贡献，请自由使用以下开发者指南 https://github.com/pytorch/pytorch/wiki/Writing-memory-format-aware-operators。

以下代码用于恢复 torch 的属性。

for (m, attrs) in old_attrs.items():
    for (k, v) in attrs.items():
        setattr(m, k, v)

待完成工作 ¶

还有许多事情要做，例如：

  • 解决 N1HW 和 NC11 张量的歧义；

  • 分布式训练支持的测试；

  • 提高操作符覆盖率。

如果您有任何反馈和建议，请通过创建问题来告诉我们。

脚本总运行时间：（0 分钟 0.000 秒）

由 Sphinx-Gallery 生成的画廊

评分这个教程

© 版权所有 2024，PyTorch。

使用 Sphinx 构建，主题由 Read the Docs 提供。
//暂时添加调查链接

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取初学者和高级开发者的深入教程

查看教程

资源

查找开发资源并获得您的疑问解答

查看资源

为了分析流量并优化您的体验，我们在本网站上提供 cookies。通过点击或导航，您同意允许我们使用 cookies。作为本站点的当前维护者，适用 Facebook 的 Cookies 政策。了解更多信息，包括可用的控制选项：Cookies 政策。