3、GPU平台部署（CUDA）：CUDA编程模型入门、cuDNN与TensorRT加速、PyTorch/TensorFlow模型导出与部署

好，咱们进入GPU部署这一章。说实话，这是整个异构计算里最硬核、也最有趣的部分。你想想看，CPU上跑得慢吞吞的模型，到了GPU上就像开了挂一样。但前提是——你得会用CUDA这把钥匙。

3.1 CUDA编程模型入门

CUDA的全称是Compute Unified Device Architecture。说白了，就是NVIDIA给开发者开的一扇后门，让你能直接指挥GPU干活。

我刚开始接触CUDA时，最困惑的就是「线程」这个概念。CPU的线程是重量级的，几百个就顶天了。但GPU的线程是轻量级的，几万个甚至几十万个同时跑都没问题。为什么会这样？因为GPU的设计哲学就是「人多力量大」。

来看一个最简单的CUDA程序——向量加法：

__global__ void vecAdd(float* A, float* B, float* C, int N) {
    int i = threadIdx.x + blockIdx.x * blockDim.x;
    if (i < N) {
        C[i] = A[i] + B[i];
    }
}

int main() {
    // 分配显存
    float *d_A, *d_B, *d_C;
    cudaMalloc(&d_A, N * sizeof(float));
    cudaMalloc(&d_B, N * sizeof(float));
    cudaMalloc(&d_C, N * sizeof(float));

    // 拷贝数据到GPU
    cudaMemcpy(d_A, h_A, N * sizeof(float), cudaMemcpyHostToDevice);
    cudaMemcpy(d_B, h_B, N * sizeof(float), cudaMemcpyHostToDevice);

    // 启动核函数
    vecAdd<<<(N+255)/256, 256>>>(d_A, d_B, d_C, N);

    // 拷贝结果回CPU
    cudaMemcpy(h_C, d_C, N * sizeof(float), cudaMemcpyDeviceToHost);

    // 释放显存
    cudaFree(d_A); cudaFree(d_B); cudaFree(d_C);
    return 0;
}

关键概念速览：

  • Grid：一个核函数启动的所有线程块集合
  • Block：一组协同工作的线程，共享一块共享内存
  • Thread：最小的执行单元，每个线程执行相同的指令

这里有个坑，我踩过好几次。就是threadIdx和blockIdx的索引计算。很多人直接写threadIdx.x就完事了，但如果你有多个block，必须加上blockIdx.x * blockDim.x。嗯，这个错误我至少犯过三次。

我的习惯：每次写CUDA代码前，先在纸上画一下线程索引的映射关系。别嫌麻烦，这能省下你两小时的调试时间。

3.2 cuDNN与TensorRT加速

手写CUDA核函数虽然很酷，但实际部署时，我们更常用的是NVIDIA的加速库。cuDNN和TensorRT就是其中的王牌。

cuDNN是CUDA深度神经网络库。它把卷积、池化、归一化这些操作都优化到了极致。我记得有一次，我把一个手写的卷积替换成cuDNN的cudnnConvolutionForward，速度直接快了8倍。你想想看，同样的硬件，换个库就快了一个数量级。

cuDNN的使用套路其实很固定：

// 创建卷积描述符
cudnnCreateConvolutionDescriptor(&convDesc);
cudnnSetConvolution2dDescriptor(convDesc, pad_h, pad_w, 
                                 stride_h, stride_w, 
                                 1, 1, CUDNN_CROSS_CORRELATION, 
                                 CUDNN_DATA_FLOAT);

// 选择算法（这一步很关键）
cudnnConvolutionFwdAlgo_t algo;
cudnnGetConvolutionForwardAlgorithm(cudnnHandle, 
                                     inputDesc, filterDesc,
                                     convDesc, outputDesc,
                                     CUDNN_CONVOLUTION_FWD_PREFER_FASTEST,
                                     0, &algo);

// 执行卷积
cudnnConvolutionForward(cudnnHandle, &alpha,
                         inputDesc, input,
                         filterDesc, filter,
                         convDesc, algo,
                         &beta, outputDesc, output);

注意：cuDNN的算法选择有四种模式：PREFER_FASTEST（最快）、PREFER_DETERMINISTIC（确定结果）、NO_WORKSPACE（省显存）、PREFER_TENSOR_CORE（用Tensor Core）。我建议你实际部署时，先用PREFER_FASTEST跑一遍，如果显存不够再换。

TensorRT就更厉害了。它不光优化算子，还能做图优化、量化、层融合。说白了，就是把你的模型「编译」成GPU最擅长执行的形式。

我曾经把一个ResNet-50模型用TensorRT优化后，推理速度从15ms降到了3ms。客户当时都惊呆了，问我是不是换了硬件。我说没有，只是把模型「翻译」成了GPU能听懂的语言。

TensorRT的核心流程：

  1. 构建期：加载模型，设置优化策略（FP32/FP16/INT8）
  2. 序列化：把优化后的引擎保存成文件（.trt或.engine）
  3. 推理期：加载引擎，执行推理

3.3 PyTorch/TensorFlow模型导出与部署

好，到了最实战的部分。怎么把训练好的模型部署到GPU上？

PyTorch模型导出：

我个人习惯用TorchScript。它能把PyTorch模型转换成一种中间表示，不依赖Python环境也能跑。

# 方法一：Tracing（推荐）
model = torch.load('model.pth')
model.eval()
example_input = torch.randn(1, 3, 224, 224)
traced_model = torch.jit.trace(model, example_input)
traced_model.save('model.pt')

# 方法二：Scripting（适合有控制流的模型）
scripted_model = torch.jit.script(model)
scripted_model.save('model_script.pt')

避坑指南：我曾经用Tracing导出一个带if-else的模型，结果推理时永远只走一个分支。后来改用Scripting才解决。所以，如果你的模型里有控制流，老老实实用torch.jit.script。

TensorFlow模型导出：

TensorFlow这边，我推荐用SavedModel格式。它把模型结构和权重打包在一起，部署时直接加载就行。

# 保存模型
model.save('my_model', save_format='tf')

# 或者用Keras的API
tf.saved_model.save(model, 'my_model')

# 加载模型
loaded = tf.saved_model.load('my_model')

部署到TensorRT：

不管是PyTorch还是TensorFlow的模型，最终都可以转到TensorRT上跑。这里有两种方式：

方式 适用场景 我的评价
ONNX转TensorRT 跨框架部署 通用性强，但ONNX算子支持有限
直接使用TensorRT API 性能极致 代码量大，但可控性高
TF-TRT（TensorFlow集成） TF用户 最省事，但优化深度有限

我个人最常用的是ONNX路线。先把PyTorch模型转成ONNX，再用TensorRT的trtexec工具转成引擎文件。流程大概是：

# 1. PyTorch转ONNX
torch.onnx.export(model, dummy_input, "model.onnx",
                  input_names=['input'],
                  output_names=['output'],
                  dynamic_axes={'input': {0: 'batch_size'},
                                'output': {0: 'batch_size'}})

# 2. ONNX转TensorRT引擎
trtexec --onnx=model.onnx --saveEngine=model.trt --fp16

核心要点：

  • 导出时一定要设置model.eval()，关掉dropout和batch norm的训练模式
  • 动态batch是个好东西，但会增加引擎构建时间。如果生产环境batch固定，建议用静态batch
  • FP16推理速度是FP32的2倍左右，但精度损失通常可以忽略

嗯，说到FP16，我想起一个案例。有一次我把一个语义分割模型转成FP16，结果发现分割边界出现了很多噪点。排查了半天，发现是某个层的数值范围太大，FP16存不下。后来在那个层前面加了个LayerNorm，问题就解决了。所以，量化不是万能的，该保留精度的层还是要保留。

最后说一句，GPU部署这件事，理论是一回事，实战是另一回事。你可能会遇到各种奇怪的问题：显存泄漏、算子不支持、精度不对齐...别慌，这些都是正常的。我刚开始做的时候，一个模型部署搞了整整两周。现在回头看，那些坑都是宝贵的经验。

下一章，咱们聊聊NPU和FPGA的部署。那个世界，又是另一番风景了。

← 上一章 📖 返回目录 下一章 →