forked from ccf-ai-infra/GPUCodeForces
Merge pull request 'finish cosfaceloss #21' (#317) from hli28146/GPUCodeForces:h21 into main
This commit is contained in:
commit
7ef720c8bf
|
|
@ -0,0 +1,202 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
cpp_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <string>
|
||||
|
||||
torch::Tensor cosface_cuda_forward(
|
||||
const torch::Tensor& cosine,
|
||||
const torch::Tensor& label,
|
||||
float s,
|
||||
float m,
|
||||
std::string reduction);
|
||||
"""
|
||||
|
||||
cuda_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <math.h>
|
||||
#include <cfloat>
|
||||
|
||||
#define BLOCK_SIZE 256
|
||||
#define WARP_SIZE 32
|
||||
|
||||
// --- Reductions Helpers ---
|
||||
template<typename T>
|
||||
__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<typename T>
|
||||
__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<<<batch_size, BLOCK_SIZE>>>(
|
||||
output.data_ptr<float>(),
|
||||
cosine.data_ptr<float>(),
|
||||
label.data_ptr<int64_t>(),
|
||||
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)
|
||||
|
|
@ -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']
|
||||
|
|
@ -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']
|
||||
|
|
@ -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()
|
||||
Loading…
Reference in New Issue