Softmax函数

概念

在数学，尤其是概率论和相关领域中，Softmax函数，或称归一化指数函数是[逻辑斯谛函数]的一种推广。它能将一个含任意实数的K维向量 之间，并且所有元素的和为1(也可视为一个 (k-1)维的hyperplane或subspace)。该函数的形式通常按下面的式子给出

实际实现的时候，为了防止溢出，会先把每个元素减去原先的最大值

实现

python实现

import numpy as np

def softmax(x):
   e_x = np.exp(x - np.max(x)) # 避免指数溢出
   return e_x / np.sum(e_x, axis=0)

cuda实现

__global__ void softmax(float *predict, int length, int *max_index) 
{
    extern __shared__ float shared_data[];
    float *shared_max_vals = shared_data;
    int *shared_max_indices = (int*)&shared_max_vals[blockDim.x];
    
    int tid = threadIdx.x;

    // 1. 找到最大值和最大值的下标，存储在共享内存中
    float max_val = -FLT_MAX;
    int max_idx = -1;
    for (int i = tid; i < length; i += blockDim.x) {
        if (predict[i] > max_val) {
            max_val = predict[i];
            max_idx = i;
        }
    }
    shared_max_vals[tid] = max_val;
    shared_max_indices[tid] = max_idx;
    __syncthreads();

    // 在所有线程间找到全局最大值和对应的下标
    if (tid == 0) {
        for (int i = 1; i < blockDim.x; i++) {
            if (shared_max_vals[i] > shared_max_vals[0]) {
                shared_max_vals[0] = shared_max_vals[i];
                shared_max_indices[0] = shared_max_indices[i];
            }
        }
        *max_index = shared_max_indices[0];
    }
    __syncthreads();

    max_val = shared_max_vals[0];

    // 2. 计算指数并求和
    float sum_exp = 0.0f;
    for (int i = tid; i < length; i += blockDim.x) {
        predict[i] = expf(predict[i] - max_val);
        sum_exp += predict[i];
    }
    shared_max_vals[tid] = sum_exp;
    __syncthreads();

    // 汇总所有线程的指数和
    if (tid == 0) {
        for (int i = 1; i < blockDim.x; i++) {
            shared_max_vals[0] += shared_max_vals[i];
        }
    }
    __syncthreads();
    float total_sum = shared_max_vals[0];

    // 3. 每个元素除以总和，得到 softmax 值
    for (int i = tid; i < length; i += blockDim.x) {
        predict[i] /= total_sum;
    }
}