From 86c6ed7d9959bc79706885c6692547f6b9008550 Mon Sep 17 00:00:00 2001 From: hli28146 Date: Tue, 2 Dec 2025 14:15:27 +0800 Subject: [PATCH] finish cosfaceloss #21 --- S1/hli28146_#21/cosfaceloss_cuda.py | 202 +++++++++++++++++++++++++++ S1/hli28146_#21/cosfaceloss_torch.py | 65 +++++++++ S1/hli28146_#21/prompt.txt | 104 ++++++++++++++ S1/hli28146_#21/run_code.py | 74 ++++++++++ 4 files changed, 445 insertions(+) create mode 100644 S1/hli28146_#21/cosfaceloss_cuda.py create mode 100644 S1/hli28146_#21/cosfaceloss_torch.py create mode 100644 S1/hli28146_#21/prompt.txt create mode 100644 S1/hli28146_#21/run_code.py diff --git a/S1/hli28146_#21/cosfaceloss_cuda.py b/S1/hli28146_#21/cosfaceloss_cuda.py new file mode 100644 index 00000000..e47156f2 --- /dev/null +++ b/S1/hli28146_#21/cosfaceloss_cuda.py @@ -0,0 +1,202 @@ +import torch +import torch.nn as nn +from torch.utils.cpp_extension import load_inline + +cpp_source = """ +#include +#include + +torch::Tensor cosface_cuda_forward( + const torch::Tensor& cosine, + const torch::Tensor& label, + float s, + float m, + std::string reduction); +""" + +cuda_source = """ +#include +#include +#include +#include + +#define BLOCK_SIZE 256 +#define WARP_SIZE 32 + +// --- Reductions Helpers --- +template +__device__ __forceinline__ T warp_reduce_max(T val) { + #pragma unroll + for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) { + val = max(val, __shfl_down_sync(0xffffffff, val, offset)); + } + return val; +} + +template +__device__ __forceinline__ T warp_reduce_sum(T val) { + #pragma unroll + for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) { + val += __shfl_down_sync(0xffffffff, val, offset); + } + return val; +} + +__device__ __forceinline__ float block_reduce_max(float val) { + static __shared__ float shared[32]; + int lane = threadIdx.x % WARP_SIZE; + int wid = threadIdx.x / WARP_SIZE; + val = warp_reduce_max(val); + if (lane == 0) shared[wid] = val; + __syncthreads(); + val = (threadIdx.x < blockDim.x / WARP_SIZE) ? shared[lane] : -FLT_MAX; + if (wid == 0) val = warp_reduce_max(val); + return val; +} + +__device__ __forceinline__ float block_reduce_sum(float val) { + static __shared__ float shared[32]; + int lane = threadIdx.x % WARP_SIZE; + int wid = threadIdx.x / WARP_SIZE; + val = warp_reduce_sum(val); + if (lane == 0) shared[wid] = val; + __syncthreads(); + val = (threadIdx.x < blockDim.x / WARP_SIZE) ? shared[lane] : 0.0f; + if (wid == 0) val = warp_reduce_sum(val); + return val; +} + +// --- Fused CosFace Kernel --- +__global__ void cosface_kernel( + float* __restrict__ output, + const float* __restrict__ cosine, + const int64_t* __restrict__ label, + int num_classes, + float s, + float m) +{ + // One block per sample (row) + int row = blockIdx.x; + int tid = threadIdx.x; + + const float* row_cosine = cosine + row * num_classes; + int64_t target_idx = label[row]; + + // 1. Pass 1: Find Max (of Modified Logits) + float local_max = -FLT_MAX; + + // Grid-Stride Loop + for (int i = tid; i < num_classes; i += blockDim.x) { + float val = row_cosine[i]; + // Apply CosFace Margin Logic On-The-Fly + if (i == target_idx) { + val = val - m; + } + // Scale + val = val * s; + + local_max = fmaxf(local_max, val); + } + + float global_max = block_reduce_max(local_max); + + // Broadcast max to all threads in block + __shared__ float s_max; + if (tid == 0) s_max = global_max; + __syncthreads(); + global_max = s_max; + + // 2. Pass 2: Sum Exp + float local_sum = 0.0f; + float target_logit = 0.0f; // Cache target logit for final loss + + for (int i = tid; i < num_classes; i += blockDim.x) { + float val = row_cosine[i]; + if (i == target_idx) { + val = val - m; + } + val = val * s; + + // Save target logit if this thread handles it + if (i == target_idx) { + target_logit = val; + } + + local_sum += __expf(val - global_max); + } + + float global_sum = block_reduce_sum(local_sum); + + // 3. Final Calculation (Thread 0) + // We need target_logit available to Thread 0. + // Optimization: Since we don't store target_logit in Shared Memory, + // let Thread 0 re-calculate it. It's just one memory read, cheaper than shared mem sync. + + if (tid == 0) { + float t_val = row_cosine[target_idx]; + t_val = (t_val - m) * s; // Target logit + + // LogSumExp = global_max + log(global_sum) + // LogSoftmax(target) = target_logit - LogSumExp + // NLL Loss = -LogSoftmax(target) = LogSumExp - target_logit + + float loss = (global_max + __logf(global_sum)) - t_val; + output[row] = loss; + } +} + +torch::Tensor cosface_cuda_forward( + const torch::Tensor& cosine, + const torch::Tensor& label, + float s, + float m, + std::string reduction) +{ + TORCH_CHECK(cosine.is_cuda() && label.is_cuda(), "Inputs must be CUDA"); + TORCH_CHECK(cosine.is_contiguous(), "Cosine must be contiguous"); + TORCH_CHECK(cosine.dim() == 2, "Cosine must be 2D (Batch, Class)"); + TORCH_CHECK(label.dim() == 1, "Label must be 1D (Batch)"); + TORCH_CHECK(cosine.size(0) == label.size(0), "Batch size mismatch"); + + int batch_size = cosine.size(0); + int num_classes = cosine.size(1); + + auto output = torch::empty({batch_size}, cosine.options()); + + // One Block per Sample + cosface_kernel<<>>( + output.data_ptr(), + cosine.data_ptr(), + label.data_ptr(), + num_classes, + s, + m + ); + + if (reduction == "mean") { + return output.mean(); + } else if (reduction == "sum") { + return output.sum(); + } + + return output; +} +""" + +class ModelNew(nn.Module): + def __init__(self, s=64.0, m=0.35, reduction='none'): + super(ModelNew, self).__init__() + self.s = s + self.m = m + self.reduction = reduction + self.op = load_inline( + name='cosface_loss_op', + cpp_sources=cpp_source, + cuda_sources=cuda_source, + functions=['cosface_cuda_forward'], + verbose=False, + extra_cuda_cflags=['-O3'] + ) + + def forward(self, cosine: torch.Tensor, label: torch.Tensor) -> torch.Tensor: + return self.op.cosface_cuda_forward(cosine.contiguous(), label.contiguous(), self.s, self.m, self.reduction) \ No newline at end of file diff --git a/S1/hli28146_#21/cosfaceloss_torch.py b/S1/hli28146_#21/cosfaceloss_torch.py new file mode 100644 index 00000000..fef7df9f --- /dev/null +++ b/S1/hli28146_#21/cosfaceloss_torch.py @@ -0,0 +1,65 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +# 人脸识别场景：Batch Size 较小，但分类数 (Classes/Identities) 极大 +BATCH_SIZE = 512 +NUM_CLASSES = 10000 +SHAPE = (BATCH_SIZE, NUM_CLASSES) + +# CosFace 超参数 +SCALE_S = 64.0 +MARGIN_M = 0.35 + +class CosFaceLoss(nn.Module): + """ + Standard PyTorch implementation of CosFace Loss. + """ + def __init__(self, s=64.0, m=0.35, reduction='mean'): + super(CosFaceLoss, self).__init__() + self.s = s + self.m = m + self.reduction = reduction + + def forward(self, cosine: torch.Tensor, label: torch.Tensor) -> torch.Tensor: + # cosine: (N, C) - Normalized Features @ Normalized Weights + # label: (N) + + # 1. 创建 One-hot 掩码 + one_hot = torch.zeros_like(cosine) + one_hot.scatter_(1, label.view(-1, 1), 1.0) + + # 2. Apply Margin: cos_theta - m (only for target class) + logits = cosine - one_hot * self.m + + # 3. Scale: s * logits + logits = logits * self.s + + # 4. CrossEntropy + loss = F.cross_entropy(logits, label, reduction='none') + + if self.reduction == 'mean': + return loss.mean() + elif self.reduction == 'sum': + return loss.sum() + return loss + +class Model(nn.Module): + def __init__(self, s=64.0, m=0.35, reduction='none'): + super(Model, self).__init__() + self.loss_fn = CosFaceLoss(s=s, m=m, reduction=reduction) + + def forward(self, cosine, label): + return self.loss_fn(cosine, label) + +def get_inputs(): + # 模拟归一化后的 Cosine Similarity (-1 ~ 1) + cosine = torch.randn(SHAPE, dtype=torch.float32) + # 归一化到 [-1, 1] 模拟真实余弦值 + cosine = torch.clamp(cosine, -1.0, 1.0) + + label = torch.randint(0, NUM_CLASSES, (BATCH_SIZE,), dtype=torch.long) + return [cosine.contiguous(), label.contiguous()] + +def get_init_inputs(): + return [SCALE_S, MARGIN_M, 'none'] \ No newline at end of file diff --git a/S1/hli28146_#21/prompt.txt b/S1/hli28146_#21/prompt.txt new file mode 100644 index 00000000..85c97db0 --- /dev/null +++ b/S1/hli28146_#21/prompt.txt @@ -0,0 +1,104 @@ +Write a custom CUDA kernel to optimize `CosFace Loss` (Large Margin Cosine Loss). + +Formula: Loss = -log( exp(s * (cos_theta_yi - m)) / Sum(exp(s * cos_theta_j_modified)) ) +Where: +- `cos_theta` is the cosine similarity matrix (Batch, Classes). +- `yi` is the ground truth class index for the sample. +- `m` is the additive cosine margin. +- `s` is the scaling factor. +- For the target class `j == yi`, the logit is `s * (cos_theta - m)`. +- For other classes `j != yi`, the logit is `s * cos_theta`. + +Problem Analysis: +1. Memory Overhead: A standard implementation uses `torch.scatter` or `one_hot` multiplication to subtract `m` only from the target indices. This allocates auxiliary tensors equal to the size of the logits (N, C), wasting memory bandwidth. +2. Operator Chaining: The sequence `scatter/sub` -> `scale` -> `CrossEntropy (LogSoftmax -> NLL)` involves multiple kernel launches and redundant global memory reads/writes. + +Optimization Strategy: Fused Logits-Modification and CrossEntropy + +The strategy is to fuse the margin application, scaling, softmax normalization, and loss calculation into a single pass. + +1. One-Block-per-Row: Each CUDA block processes one sample (one row of the cosine matrix) to calculate its loss. + +2. On-the-Fly Logic: + - Load cosine values from global memory. + - Check if the current column index matches the target label `y_i`. + - If match: val = s * (val - m). + - If not match: val = s * val. + This eliminates the need for one-hot masks or scatter operations. + +3. Online Softmax (LogSumExp): + - Use the 2-pass reduction algorithm (or 1-pass online algorithm) within the block to compute `Max` and `SumExp` of the *modified* logits. + - This ensures numerical stability without materializing the full modified logits tensor. + +4. Fused NLL Loss: + - Calculate `log_prob = target_logit - (max_val + log(sum_exp))`. + - Output `-log_prob`. + +Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is: + +```python +import torch +import torch.nn as nn +import torch.nn.functional as F + +# 人脸识别场景：Batch Size 较小，但分类数 (Classes/Identities) 极大 +BATCH_SIZE = 512 +NUM_CLASSES = 10000 +SHAPE = (BATCH_SIZE, NUM_CLASSES) + +# CosFace 超参数 +SCALE_S = 64.0 +MARGIN_M = 0.35 + +class CosFaceLoss(nn.Module): + """ + Standard PyTorch implementation of CosFace Loss. + """ + def __init__(self, s=64.0, m=0.35, reduction='mean'): + super(CosFaceLoss, self).__init__() + self.s = s + self.m = m + self.reduction = reduction + + def forward(self, cosine: torch.Tensor, label: torch.Tensor) -> torch.Tensor: + # cosine: (N, C) - Normalized Features @ Normalized Weights + # label: (N) + + # 1. 创建 One-hot 掩码 + one_hot = torch.zeros_like(cosine) + one_hot.scatter_(1, label.view(-1, 1), 1.0) + + # 2. Apply Margin: cos_theta - m (only for target class) + logits = cosine - one_hot * self.m + + # 3. Scale: s * logits + logits = logits * self.s + + # 4. CrossEntropy + loss = F.cross_entropy(logits, label, reduction='none') + + if self.reduction == 'mean': + return loss.mean() + elif self.reduction == 'sum': + return loss.sum() + return loss + +class Model(nn.Module): + def __init__(self, s=64.0, m=0.35, reduction='none'): + super(Model, self).__init__() + self.loss_fn = CosFaceLoss(s=s, m=m, reduction=reduction) + + def forward(self, cosine, label): + return self.loss_fn(cosine, label) + +def get_inputs(): + # 模拟归一化后的 Cosine Similarity (-1 ~ 1) + cosine = torch.randn(SHAPE, dtype=torch.float32) + # 归一化到 [-1, 1] 模拟真实余弦值 + cosine = torch.clamp(cosine, -1.0, 1.0) + + label = torch.randint(0, NUM_CLASSES, (BATCH_SIZE,), dtype=torch.long) + return [cosine.contiguous(), label.contiguous()] + +def get_init_inputs(): + return [SCALE_S, MARGIN_M, 'none'] \ No newline at end of file diff --git a/S1/hli28146_#21/run_code.py b/S1/hli28146_#21/run_code.py new file mode 100644 index 00000000..14819cd5 --- /dev/null +++ b/S1/hli28146_#21/run_code.py @@ -0,0 +1,74 @@ +########################################################### +# 性能和精度验证程序 +########################################################### +import torch +import torch.nn as nn +import time +from cosfaceloss_torch import Model,get_inputs,get_init_inputs +from cosfaceloss_cuda import ModelNew + +def run_benchmark(): + # 检查 CUDA 是否可用 + if not torch.cuda.is_available(): + print("CUDA 不可用，请确保您有可用的 NVIDIA GPU 并已正确安装 PyTorch CUDA 版本。") + return + else: + device = torch.device("cuda") + + # 初始化模型 + init_inputs = get_init_inputs() + init_inputs = [ + x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in init_inputs + ] + inputs = get_inputs() + inputs = [ + x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in inputs + ] + + torch_model = Model(*init_inputs).cuda() + cuda_model = ModelNew(*init_inputs).cuda() + + torch_model.eval() + cuda_model.eval() + + print("-------------------- 精度对齐验证 --------------------") + with torch.no_grad(): + output_torch = torch_model( *inputs) + output_cuda = cuda_model(*inputs) + + precision_flag = torch.allclose(output_torch, output_cuda,rtol=1e-03) + if precision_flag: + print("✅ 精度对齐：两个模型的输出结果非常接近。") + else: + print("❌ 精度不一致！") + + print("\n-------------------- 性能加速比测试 --------------------") + num_iterations = 100 + + # PyTorch 模型计时 + torch.cuda.synchronize() + start_time = time.time() + for _ in range(num_iterations): + _ = torch_model(*inputs) + torch.cuda.synchronize() + torch_time = (time.time() - start_time) / num_iterations + + # 自定义 CUDA 内核计时 + torch.cuda.synchronize() + start_time = time.time() + for _ in range(num_iterations): + _ = cuda_model(*inputs) + torch.cuda.synchronize() + cuda_time = (time.time() - start_time) / num_iterations + + print(f"PyTorch torch.relu 平均执行时间: {torch_time:.6f} 秒") + print(f"自定义 CUDA 内核 平均执行时间: {cuda_time:.6f} 秒") + speedup = 0 + if cuda_time > 0: + speedup = torch_time / cuda_time + print(f"加速比 (Speedup): {speedup:.2f}x") + else: + print("CUDA 内核执行时间为0，无法计算加速比。") + return precision_flag,speedup +if __name__ == "__main__": + precision_flag,speedup = run_benchmark() \ No newline at end of file