Caffe、Tensorflow和Pytorch通道维度顺序小结

Posted on | In |

Caffe、Tensorflow和Pytorch通道维度顺序小结

N: batch;

C: channel

H: height

W: width

Caffe 的Blob通道顺序是：NCHW;

Tensorflow的tensor通道顺序：默认是NHWC， 也支持NCHW，使用cuDNN会更快;

Pytorch中tensor的通道顺序：NCHW

TensorRT中的tensor 通道顺序： NCHW

numpy image: H x W x C
torch image: C X H X W

目前世面上深度学习框架比较多，常用的有三大类：Caffe、Tensorflow和Pytorch，这三种深度学习框架都主要在英伟达显卡上面进行训练和测试，很奇怪的是，它们之间的通道维度顺序并没有保持一致，在多个框架中替换着使用，很容易让人混淆了，所以做一个小结。

Caffe 的通道顺序是NCHW;

Tensorflow的通道顺序默认是NHWC（但可以设置成NCHW），NHWC 的访存局部性更好（每三个输入像素即可得到一个输出像素），NCHW 则必须等所有通道输入准备好才能得到最终输出结果，需要占用较大的临时空间。

TensorFlow 为什么选择 NHWC 格式作为默认格式？因为早期开发都是基于 CPU，使用 NHWC 比 NCHW 稍快一些（不难理解，NHWC 局部性更好，cache 利用率高）。

NCHW 则是 Nvidia cuDNN 默认格式，使用 GPU 加速时用 NCHW 格式速度会更快（也有个别情况例外）。

最佳实践：设计网络时充分考虑两种格式，最好能灵活切换，在 GPU 上训练时使用 NCHW 格式，在 CPU 上做预测时使用 NHWC 格式。

Pytorch的通道顺序是NCHW

下采样

常见下采样方法
最大池化
平均池化
卷积（strides>1）
空洞卷积(Atrous Convolution / Dilated Convolution) 也是一种变相的下采样方式，但它通常用于保持特征图尺寸的同时增加感受野
自适应下采样(Adaptive Sampling)：

上采样

（将低分辨率特征图放大至高分辨率）

Option Function
Nearest Neighbor Upsampling
（最近邻插值
Bilinear Interpolation Upsampling
（双线性插值）		 old函数 nn.Upsample()
F.interpolate (1.6版本之后)
Transposed Convolution（转置卷积/反卷积） nn.ConvTranspose2d
PixelShuffle（像素重组） nn.PixelShuffle

spectral_norm

是深度学习中的一种正则化技术，主要用于稳定和约束神经网络中的权重矩阵，特别是对于生成对抗网络（GANs）和其他涉及大型权重矩阵的模型，可以有效地解决训练过程中的梯度消失或爆炸问题，以及防止模型过度拟合。

Spectral normalization（谱范数归一化）通过对权重矩阵的谱范数（即该矩阵的最大奇异值）进行规范化，强制限制权重矩阵的影响范围，从而使网络的训练更加稳定。

  1. 给定一个卷积层 conv，我们可以应用谱范数归一化：
1
2
3
4
5
6
7
8
from torch.nn.utils import spectral_norm
from torch import nn

# 创建一个卷积层
conv = nn.Conv2d(in_channels, out_channels, kernel_size, bias=False)

# 对权重矩阵应用谱范数归一化
conv = spectral_norm(conv)
  1. 或者在初始化时直接使用 nn.SpectralNorm：
1
2
3
4
5
from torch.nn import SpectralNorm

# 创建并直接应用谱范数归一化的卷积层
conv = nn.Conv2d(in_channels, out_channels, kernel_size, bias=False)
conv = SpectralNorm(conv)

Container Class

nn.Sequential

是一个有序的模块容器，其中的子模块会按照添加的顺序依次执行。它非常适用于堆叠一系列简单的线性操作序列，如连续的卷积和全连接层。定义和使用方式如下：

1
2
3
4
5
6
7
8
from torch.nn import Sequential, Conv2d, Linear

model = Sequential(
    Conv2d(in_channels=3, out_channels=64, kernel_size=3, padding=1),
    ReLU(),
    MaxPool2d(kernel_size=2),
    Linear(64 * reduced_image_size * reduced_image_size, num_classes)
)

nn.ModuleList

是一个可迭代的模块容器，但不保持顺序索引。它可以容纳任意数量的 nn.Module 子类实例，并且可以通过索引访问和修改。

1
2
3
4
from torch.nn import ModuleList, Conv2d
conv_layers = ModuleList([Conv2d(32, 64, 3), Conv2d(64, 128, 3)])

self_layers = ModuleList([MyModule(12, 24), MyModule2(24,12)])

nn.ModuleDict

nn.ModuleDict 是一个键值对形式的模块容器，其中键是字符串类型，值是 nn.Module 实例。它允许通过字符串关键字来访问和管理子模块，这对于具有命名组件的复杂网络结构非常有用。

1
2
3
4
5
6
from torch.nn import ModuleDict, Conv2d

conv_blocks = ModuleDict({
    'block1': Conv2d(3, 32, 3),
    'block2': Conv2d(32, 64, 5),
})

functional

Leak_relu

F.leaky_relu_() 和 F.leaky_relu() 都是用来实现Leaky ReLU激活函数的，它们的区别在于是否原地修改输入张量：

  1. F.leaky_relu()： 这是一个普通的函数调用，它接收一个张量作为输入，计算并返回带有Leaky ReLU激活的输出张量。这个操作不会改变原始输入张量的内容，而是返回一个新的张量。

    1
    2
    3
    1import torch.nn.functional as F
    2input = torch.randn(10, 10)
    3output = F.leaky_relu(input, negative_slope=0.01)  # 返回一个新的张量

  2. F.leaky_relu_()： 这是一个原地(in-place)操作符版本，它会在原始输入张量上直接进行Leaky ReLU激活计算，覆盖掉原来的值，不返回新的张量，而是直接修改输入张量。

    1
    2
    3
    1import torch.nn.functional as F
    2input = torch.randn(10, 10)
    3F.leaky_relu_(input, negative_slope=0.01)  # 直接在原始输入上修改

总结来说，如果你希望保留原始输入张量以便后续使用，请使用 F.leaky_relu()；如果你愿意直接在原始输入上进行操作且不需要保留原始值，可以使用 F.leaky_relu_()，这在内存有限的情况下有助于节省存储空间。