为什么我的卷积实现比 Tensorflow 的慢？

Question

我已经在 C++ 中使用 SIMD 指令实现了 VGG19 网络，仅用于推理。我想优化一个推理请求的延迟。

由于 VGG19 主要由卷积层组成，因此我主要专注于实现高效的卷积层。我在做的时候跟着这篇论文：Anatomy Of High-Performance Deep Learning Convolutions On SIMD Architectures.

我的实施提供了正确的结果。我使用 SIMD Intrisics 和论文中描述的算法。所有的重量都是预先加载的。每层的输入和输出缓冲区是在运行实际推理之前分配的。

---

我们以 VGG19 网络的第二个卷积层为例：

• 输入：（224、224、64）（填充后的226、226、64）
• 输出：(224, 224, 64)
• 内核：(3, 3, 64, 64) (KH, KW, C_IN, C_OUT)

下面是代码对应的代码：

void conv2d_block1_conv2(const float* in, const float* weights, float* out) {
    constexpr int VLEN = 8; // to use _mm256_* intrisics
    constexpr int C_OUT_B = VLEN;
    constexpr int C_IN_B = VLEN;

    constexpr int H = 226;           // Input Height
    constexpr int W = 226;           // Input Width
    constexpr int C_IN = 64;         // Input Channels

    constexpr int KH = 3;            // Kernel Height
    constexpr int KW = 3;            // Kernel Width

    constexpr int H_OUT = 224;       // Output Height
    constexpr int W_OUT = 224;       // Output Width
    constexpr int C_OUT = 64;        // Output Channels

    __m256 in_vec, weights_vec, out_vec;
    for (int c_out = 0; c_out < C_OUT / C_OUT_B; c_out++)
    for (int c_in_b = 0; c_in_b < C_IN / C_IN_B; c_in_b++)
    for (int h_out = 0; h_out < H_OUT; h_out++)
    for (int w_out = 0; w_out < W_OUT; w_out++){
        const int outIdx = LINEAR_4(c_out, h_out, w_out, 0, H_OUT, W_OUT, C_OUT_B);
        out_vec = _mm256_load_ps (&out[outIdx]);
        for (int kh = 0; kh < KH; kh++)
            for (int kw = 0; kw < KW; kw++)
                for (int c_in = 0; c_in < C_IN_B; c_in++){
                    const int inIdx = LINEAR_4(c_in_b, h_out + kh, w_out + kw, c_in, H, W, C_IN_B);
                    const int weightsIdx = LINEAR_6(c_out, c_in_b, kh, kw, c_in, 0, C_IN / C_IN_B, KH, KW, C_IN_B, C_OUT_B);
                    in_vec = _mm256_set1_ps (in[inIdx]);
                    weights_vec = _mm256_load_ps(&weights[weightsIdx]); 
                    out_vec = _mm256_fmadd_ps (in_vec, weights_vec, out_vec);
                    _mm256_store_ps(&out[outIdx], out_vec);
                }
    }
}

注意：我正在研究线性地址空间。函数LINEAR4 和LINEAR6 将多维索引映射到一维索引。

array[c_out][h_out][w_out][0]         <-> LINEAR_4(c_out, h_out, w_out, 0, H_OUT, W_OUT, C_OUT_B); 
array[c_out][c_in_b][kh][kw][c_in][0] <-> LINEAR_6(c_out, c_in_b, kh, kw, c_in, 0, C_IN / C_IN_B, KH, KW, C_IN_B, C_OUT_B);

---

我为每个卷积层创建了一个类似上述的函数，为编译器提供最佳优化可能性。

但是，执行时间相当糟糕。 对于整个 VGG19 网络（都是单线程执行）：

• 我的实现：2400ms
• 带有 Tensorflow 后端的 Keras 使用 model.predict(image)：600 毫秒

这种巨大的性能差距让我想知道我做错了什么。我正在使用带有-O3 标志的clang。

所以我的问题是：

1. 是否有我没有考虑到的关键因素？
2. Keras/TensorFlow 使用的是哪个实现。他们怎么这么快？

Answer 1

我找到了性能不佳的原因。 clang 编译器只使用了 2 个 SSE 寄存器，而不是所有可用的。这导致对 L1 缓存进行不必要的写入和读取。

我手动展开了两个内部循环，编译器现在使用所有可用的 16 个 SSE 寄存器。性能大幅提升。

如果您使用 SSE Intrisics，请务必检查生成的程序集。