从零实现Transformer模型：C++实战指南与性能优化技巧

1次阅读

共计 2063 个字符，预计需要花费 6 分钟才能阅读完成。

Transformer 在机器翻译、文本生成等 NLP 任务中表现出色，但其 Python 实现常因动态类型和解释执行导致性能瓶颈。C++ 的静态编译特性更适合生产环境部署，但需要处理手动内存管理、缺乏现成算子库等问题。本文将带您避开这些坑，实现一个工业级可用的 Transformer 模型。

Eigen：
优点：纯头文件库，模板元编程优化，适合小矩阵运算
缺点：大矩阵乘法性能不如专用 BLAS 实现
OpenBLAS：
优点：针对 CPU 架构优化的 Level 3 BLAS
缺点：需要单独编译安装
混合方案：
小矩阵用 Eigen，大矩阵乘法调用 BLAS 的cblas_sgemm

/**
 * @brief 多头注意力计算
 * @param Q 查询矩阵 [batch_size, num_heads, seq_len, head_dim]
 * @param K 键矩阵   [batch_size, num_heads, seq_len, head_dim]
 * @param V 值矩阵   [batch_size, num_heads, seq_len, head_dim]
 * @param mask 注意力掩码
 * @note 使用预分配内存避免重复申请
 */
Eigen::Tensor<float, 4> MultiHeadAttention(
    const Eigen::Tensor<float, 4>& Q,
    const Eigen::Tensor<float, 4>& K,
    const Eigen::Tensor<float, 4>& V,
    const Eigen::Tensor<bool, 2>& mask) {// 计算 attention 分数 (Q*K^T)/sqrt(d_k)
  auto scores = Q.contract(K, 
    Eigen::array<Eigen::IndexPair<int>, 1>{{2, 2} }) 
    / std::sqrt(Q.dimension(3));

  // 应用 mask
  if (mask.dimension(0) != 0) {scores = scores + mask.cast<float>().log().reshape({1, 1, mask.dimension(0), mask.dimension(1)});
  }

  // Softmax 归一化
  auto attn = softmax(scores);

  // 加权求和
  return attn.contract(V, 
    Eigen::array<Eigen::IndexPair<int>, 1>{{3, 2} });
}

// 添加 epsilon 防止除零
constexpr float eps = 1e-5f;

Eigen::VectorXf LayerNorm(const Eigen::VectorXf& x) {float mean = x.mean();
  float var = (x.array() - mean).square().mean();
  return (x.array() - mean) / std::sqrt(var + eps);
}

// 使用 AVX2 指令集加速
#ifdef __AVX2__
#include <immintrin.h>

void matmul_avx2(const float* A, const float* B, float* C, 
                 int m, int n, int k) {for (int i = 0; i < m; ++i) {for (int j = 0; j < n; j += 8) {__m256 c = _mm256_setzero_ps();
      for (int l = 0; l < k; ++l) {__m256 a = _mm256_set1_ps(A[i*k + l]);
        __m256 b = _mm256_loadu_ps(&B[l*n + j]);
        c = _mm256_fmadd_ps(a, b, c);
      }
      _mm256_storeu_ps(&C[i*n + j], c);
    }
  }
}
#endif