本文介绍了针对InsightFace模型中PReLU激活函数的CUDA优化，特别是针对Broadcast模式的优化。通过Pack向量化的方法减少了整数除法和取余的计算，从而提高了Kernel的计算效率和带宽利用率。在A100 GPU上，优化后的PReLU实现相比于朴素实现，计算占比降低20%-30%，吞吐提升30+%，在多种形状测试中展现出接近两倍于PyTorch的性能。
撰文｜郑泽康
InsightFace模型里大量使用了PReLU激活函数，而PReLU的工作模式有两种：
1. PReLU(1)，此时权重alpha的形状为(1, )，等价于一个Elementwise操作。
2. PReLU(channels)，此时权重alpha的形状为(channels, )，和输入特征(N, C, H, W)中C的大小是对应的。此时PReLU等价于一个Binary Broadcast操作。
InsightFace模型里的PReLU工作模式是第二种，之前已经介绍过CUDA Elementwise操作优化，而在Broadcast情形下也存在一定的优化机会。
1
朴素实现
一个朴素实现的思想就是在循环内部，根据当前元素的索引，推算出该元素对应需要使用的alpha权重的索引。然后判断当前元素x是否大于0，若大于0则返回x，小于0则返回alpha*x。对应代码如下：
template<typename T>
__global__ void PReluForwardGpu(const int32_t elem_cnt, const int32_t alpha_size,
const int32_t inner_size, const T* x, const T* alpha, T* y) {
CUDA_1D_KERNEL_LOOP(i, elem_cnt) {
const T x_i = x[i];
&
最低0.47元/天 解锁文章
扫一扫
671
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
