C语言实现torch.nn.functional.multi_head_attention_forward

下面是一个使用C语言实现的torch.nn.functional.multi_head_attention_forward的示例代码：

#include <stdio.h>
#include <stdlib.h>
#include <math.h>

#define MAX_SEQ_LENGTH 256
#define MAX_HIDDEN_SIZE 512

void multi_head_attention_forward(float *input, float *weight_q, float *weight_k, float *weight_v, float *weight_o, float *bias_q, float *bias_k, float *bias_v, float *bias_o, float *output, int batch_size, int seq_length, int num_heads, int head_size, float dropout_prob) {
    float q[MAX_SEQ_LENGTH][MAX_HIDDEN_SIZE];
    float k[MAX_SEQ_LENGTH][MAX_HIDDEN_SIZE];
    float v[MAX_SEQ_LENGTH][MAX_HIDDEN_SIZE];
    float qk[MAX_SEQ_LENGTH][MAX_SEQ_LENGTH];
    float qkv[MAX_SEQ_LENGTH][MAX_HIDDEN_SIZE];
    float o[MAX_SEQ_LENGTH][MAX_HIDDEN_SIZE];
    float attention_probs[MAX_SEQ_LENGTH][MAX_SEQ_LENGTH];
    float q_bias[MAX_SEQ_LENGTH];
    float k_bias[MAX_SEQ_LENGTH];
    float v_bias[MAX_SEQ_LENGTH];
    float o_bias[MAX_SEQ_LENGTH];
    float q_scale_factor = sqrtf((float)head_size);
    float k_scale_factor = sqrtf((float)head_size);
    float v_scale_factor = sqrtf((float)head_size);
    float attention_scale_factor = 1.0f / sqrtf((float)head_size);

    // Compute q, k, v
    for (int i = 0; i < batch_size; i++) {
        for (int j = 0; j < seq_length; j++) {
            for (int h = 0; h < num_heads; h++) {
                for (int d = 0; d < head_size; d++) {
                    int input_idx = i * seq_length * num_heads * head_size + j * num_heads * head_size + h * head_size + d;
                    int q_idx = i * seq_length * num_heads * head_size + j * num_heads * head_size + h * head_size + d;
                    int k_idx = i * seq_length * num_heads * head_size + j * num_heads * head_size + h * head_size + d;
                    int v_idx = i * seq_length * num_heads * head_size + j * num_heads * head_size + h * head_size + d;

                    q[q_idx] = input[input_idx] * q_scale_factor + bias_q[h * head_size + d];
                    k[k_idx] = input[input_idx] * k_scale_factor + bias_k[h * head_size + d];
                    v[v_idx] = input[input_idx] * v_scale_factor + bias_v[h * head_size + d];
                }
            }
        }
    }

    // Compute qk
    for (int i = 0; i < batch_size * seq_length * num_heads; i++) {
        for (int j = 0; j < seq_length * num_heads; j++) {
            qk[i][j] = 0.0f;
            for (int d = 0; d < head_size; d++) {
                int q_idx = i * head_size + d;
                int k_idx = j * head_size + d;
                qk[i][j] += q[q_idx] * k[k_idx];
            }
        }
    }

    // Compute attention_probs
    for (int i = 0; i < batch_size * seq_length * num_heads; i++) {
        for (int j = 0; j < seq_length * num_heads; j++) {
            attention_probs[i][j] = expf(qk[i][j] * attention_scale_factor);
        }
    }

    // Apply dropout
    for (int i = 0; i < batch_size * seq_length * num_heads; i++) {
        for (int j = 0; j < seq_length * num_heads; j++) {
            if ((float)rand() / RAND_MAX < dropout_prob) {
                attention_probs[i][j] = 0.0f;
            }
        }
    }

    // Normalize attention_probs
    for (int i = 0; i < batch_size * seq_length * num_heads; i++) {
        float sum = 0.0f;
        for (int j = 0; j < seq_length * num_heads; j++) {
            sum += attention_probs[i][j];
        }
        for (int j = 0; j < seq_length * num_heads; j++) {
            attention_probs[i][j] /= sum;
        }
    }

    // Compute qkv
    for (int i = 0; i < batch_size * seq_length * num_heads; i++) {
        for (int j = 0; j < head_size; j++) {
            qkv[i][j] = 0.0f;
            for (int k = 0; k < seq_length * num_heads; k++) {
                int q_idx = i * head_size + j;
                int v_idx = k * head_size + j;
                qkv[i][j] += attention_probs[i][k] * v[v_idx];
            }
        }
    }

    // Compute o
    for (int i = 0; i < batch_size; i++) {
        for (int j = 0; j < seq_length; j++) {
            for (int h = 0; h < num_heads; h++) {
                for (int d = 0; d < head_size; d++) {
                    int o_idx = i * seq_length * num_heads * head_size + j * num_heads * head_size + h * head_size + d;
                    int qkv_idx = i * seq_length * num_heads * head_size + j * num_heads * head_size + h * head_size + d;

                    o[o_idx] = qkv[qkv_idx] + bias_o[h * head_size + d];
                }
            }
        }
    }

    // Compute output
    for (int i = 0; i < batch_size; i++) {
        for (int j = 0; j < seq_length; j++) {
            for (int h = 0; h < num_heads * head_size; h++) {
                int output_idx = i * seq_length * num_heads * head_size + j * num_heads * head_size + h;
                int o_idx = i * seq_length * num_heads * head_size + j * num_heads * head_size + h;

                output[output_idx] = o[o_idx];
            }
        }
    }
}

这是一个简单的多头注意力机制的前向传播函数，输入参数包括输入张量(input)、查询权重矩阵(weight_q)、键权重矩阵(weight_k)、值权重矩阵(weight_v)、输出权重矩阵(weight_o)、查询偏置向量(bias_q)、键偏置向量(bias_k)、值偏置向量(bias_v)、输出偏置向量(bias_o)、输出张量(output)、批次大小(batch_size)、序列长度(seq_length)、头数(num_heads)、头大小(head_size)和dropout概率(dropout_prob)。