Softmax函數(shù)

概念

在數(shù)學(xué)，尤其是概率論和相關(guān)領(lǐng)域中，Softmax函數(shù)，或稱歸一化指數(shù)函數(shù)是[邏輯斯諦函數(shù)]的一種推廣。它能將一個含任意實(shí)數(shù)的K維向量 之間，并且所有元素的和為1(也可視為一個 (k-1)維的hyperplane或subspace)。該函數(shù)的形式通常按下面的式子給出

實(shí)際實(shí)現(xiàn)的時候，為了防止溢出，會先把每個元素減去原先的最大值

實(shí)現(xiàn)

python實(shí)現(xiàn)

import numpy as np

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

cuda實(shí)現(xiàn)

__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. 找到最大值和最大值的下標(biāo)，存儲在共享內(nèi)存中
    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();

    // 在所有線程間找到全局最大值和對應(yīng)的下標(biāo)
    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. 計(jì)算指數(shù)并求和
    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();

    // 匯總所有線程的指數(shù)和
    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;
    }
}