forked from ccf-ai-infra/GPUCodeForces
Merge pull request 'FEAT:ADD focalloss_sigmoid #107' (#1038) from wut0n/GPUCodeForces:focalloss_sigmoid into main
This commit is contained in:
commit
5b61ff36f7
|
|
@ -0,0 +1,133 @@
|
|||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
import os
|
||||
|
||||
focal_loss_sigmoid_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <cmath>
|
||||
|
||||
__global__ void focal_loss_sigmoid_fused_kernel(
|
||||
const float* logits,
|
||||
const float* targets,
|
||||
float* loss,
|
||||
int batch_size,
|
||||
float alpha,
|
||||
float gamma
|
||||
) {
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx < batch_size) {
|
||||
float logit = logits[idx];
|
||||
float target = targets[idx];
|
||||
|
||||
// Optimized sigmoid computation with numerical stability
|
||||
// Avoids redundant computation and improves precision
|
||||
float prob;
|
||||
if (logit > 0.0f) {
|
||||
float exp_neg_logit = expf(-logit);
|
||||
prob = 1.0f / (1.0f + exp_neg_logit);
|
||||
} else {
|
||||
float exp_logit = expf(logit);
|
||||
prob = exp_logit / (1.0f + exp_logit);
|
||||
}
|
||||
|
||||
// Clamp probability to avoid log(0) - more aggressive clamping for stability
|
||||
prob = fmaxf(fminf(prob, 1.0f - 1e-8f), 1e-8f);
|
||||
|
||||
// Compute pt based on target
|
||||
float pt = (target > 0.5f) ? prob : (1.0f - prob);
|
||||
|
||||
// Compute focal weight with optimized powf call
|
||||
float focal_weight = alpha * powf(1.0f - pt, gamma);
|
||||
|
||||
// Optimized binary cross entropy computation
|
||||
// Directly compute BCE without intermediate steps
|
||||
float bce;
|
||||
if (target > 0.5f) {
|
||||
bce = -logf(prob);
|
||||
} else {
|
||||
bce = -logf(1.0f - prob);
|
||||
}
|
||||
|
||||
// Apply focal weight
|
||||
loss[idx] = focal_weight * bce;
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor focal_loss_sigmoid_fused_cuda(torch::Tensor logits, torch::Tensor targets, float alpha, float gamma) {
|
||||
// Ensure consistent data types
|
||||
logits = logits.to(torch::kFloat32);
|
||||
targets = targets.to(torch::kFloat32);
|
||||
|
||||
auto batch_size = logits.numel();
|
||||
auto loss = torch::empty_like(logits);
|
||||
|
||||
// Optimized block configuration
|
||||
const int block_size = 256;
|
||||
int num_blocks = (batch_size + block_size - 1) / block_size;
|
||||
|
||||
// Launch optimized kernel
|
||||
focal_loss_sigmoid_fused_kernel<<<num_blocks, block_size>>>(
|
||||
logits.data_ptr<float>(),
|
||||
targets.data_ptr<float>(),
|
||||
loss.data_ptr<float>(),
|
||||
batch_size,
|
||||
alpha,
|
||||
gamma
|
||||
);
|
||||
|
||||
// Check for CUDA errors
|
||||
cudaError_t err = cudaGetLastError();
|
||||
if (err != cudaSuccess) {
|
||||
AT_ERROR("CUDA kernel failed: ", cudaGetErrorString(err));
|
||||
}
|
||||
|
||||
return loss;
|
||||
}
|
||||
"""
|
||||
|
||||
focal_loss_sigmoid_cpp_source = """
|
||||
torch::Tensor focal_loss_sigmoid_fused_cuda(torch::Tensor logits, torch::Tensor targets, float alpha, float gamma);
|
||||
"""
|
||||
|
||||
build_dir = './cuda_build'
|
||||
os.makedirs(build_dir, exist_ok=True)
|
||||
|
||||
# Compile the optimized inline CUDA code
|
||||
focal_loss_sigmoid_module = load_inline(
|
||||
name="focal_loss_sigmoid_fused",
|
||||
cpp_sources=focal_loss_sigmoid_cpp_source,
|
||||
cuda_sources=focal_loss_sigmoid_source,
|
||||
functions=["focal_loss_sigmoid_fused_cuda"],
|
||||
verbose=True,
|
||||
build_directory=build_dir,
|
||||
extra_cuda_cflags=["-O3", "--use_fast_math"] # Optimized compilation flags
|
||||
)
|
||||
|
||||
class ModelNew(torch.nn.Module):
|
||||
def __init__(self, alpha=0.25, gamma=2.0, reduction='mean'):
|
||||
super(ModelNew, self).__init__()
|
||||
self.alpha = alpha
|
||||
self.gamma = gamma
|
||||
self.reduction = reduction
|
||||
self.focal_loss_sigmoid = focal_loss_sigmoid_module
|
||||
|
||||
def forward(self, logits, targets):
|
||||
# Ensure consistent input types
|
||||
logits = logits.to(torch.float32)
|
||||
targets = targets.to(torch.float32)
|
||||
|
||||
# Handle shape matching
|
||||
if logits.dim() == 2 and logits.size(1) == 1:
|
||||
logits = logits.squeeze(1)
|
||||
|
||||
# Compute focal loss using optimized fused CUDA kernel
|
||||
loss = self.focal_loss_sigmoid.focal_loss_sigmoid_fused_cuda(logits, targets, self.alpha, self.gamma)
|
||||
|
||||
# Apply reduction
|
||||
if self.reduction == 'mean':
|
||||
return loss.mean()
|
||||
elif self.reduction == 'sum':
|
||||
return loss.sum()
|
||||
else:
|
||||
return loss
|
||||
|
|
@ -0,0 +1,72 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
Focal Loss with Sigmoid implementation for binary classification.
|
||||
Fully fused version that computes sigmoid and focal loss in a single pass.
|
||||
Focal Loss = -α * (1-pt)^γ * log(pt)
|
||||
where pt = p if target=1, else (1-p), p = sigmoid(logit)
|
||||
"""
|
||||
def __init__(self, alpha=0.25, gamma=2.0, reduction='mean'):
|
||||
super(Model, self).__init__()
|
||||
self.alpha = alpha
|
||||
self.gamma = gamma
|
||||
self.reduction = reduction
|
||||
|
||||
def forward(self, inputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Compute Focal Loss with fused sigmoid computation.
|
||||
|
||||
Args:
|
||||
inputs (torch.Tensor): Predicted logits of shape (batch_size, num_classes)
|
||||
targets (torch.Tensor): Ground truth labels of shape (batch_size,)
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Computed focal loss
|
||||
"""
|
||||
# Ensure input types are consistent
|
||||
inputs = inputs.to(torch.float32)
|
||||
targets = targets.to(torch.float32)
|
||||
|
||||
# Handle shape matching
|
||||
if inputs.dim() == 2 and inputs.size(1) == 1:
|
||||
inputs = inputs.squeeze(1)
|
||||
|
||||
# Fully fused computation using highly optimized built-in functions
|
||||
# Compute sigmoid with numerical stability
|
||||
sigmoid_inputs = torch.sigmoid(inputs)
|
||||
|
||||
# Compute pt based on targets
|
||||
pt = torch.where(targets == 1, sigmoid_inputs, 1 - sigmoid_inputs)
|
||||
|
||||
# Compute focal weight
|
||||
focal_weight = self.alpha * torch.pow(1 - pt, self.gamma)
|
||||
|
||||
# Compute binary cross entropy with logits (more numerically stable)
|
||||
bce = F.binary_cross_entropy_with_logits(inputs, targets, reduction='none')
|
||||
|
||||
# Apply focal weight
|
||||
focal_loss = focal_weight * bce
|
||||
|
||||
# Apply reduction
|
||||
if self.reduction == 'mean':
|
||||
return focal_loss.mean()
|
||||
elif self.reduction == 'sum':
|
||||
return focal_loss.sum()
|
||||
else:
|
||||
return focal_loss
|
||||
|
||||
batch_size = 32
|
||||
num_classes = 1 # Binary classification
|
||||
|
||||
def get_inputs():
|
||||
# Generate random logits with explicit float32
|
||||
inputs = torch.randn(batch_size, num_classes, dtype=torch.float32)
|
||||
# Generate random binary targets (0 or 1) with explicit float32
|
||||
targets = torch.randint(0, 2, (batch_size,), dtype=torch.float32)
|
||||
return [inputs, targets]
|
||||
|
||||
def get_init_inputs():
|
||||
return [] # No special initialization inputs needed
|
||||
|
|
@ -0,0 +1,103 @@
|
|||
Write a custom CUDA kernel to replace PyTorch's Focal Loss with Sigmoid implementation for binary classification.
|
||||
|
||||
You are given the following PyTorch architecture:
|
||||
|
||||
python
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
Focal Loss with Sigmoid implementation for binary classification.
|
||||
Fully fused version that computes sigmoid and focal loss in a single pass.
|
||||
Focal Loss = -α * (1-pt)^γ * log(pt)
|
||||
where pt = p if target=1, else (1-p), p = sigmoid(logit)
|
||||
"""
|
||||
def init(self, alpha=0.25, gamma=2.0, reduction='mean'):
|
||||
super(Model, self).init()
|
||||
self.alpha = alpha
|
||||
self.gamma = gamma
|
||||
self.reduction = reduction
|
||||
|
||||
def forward(self, inputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Compute Focal Loss with fused sigmoid computation.
|
||||
|
||||
Args:
|
||||
inputs (torch.Tensor): Predicted logits of shape (batch_size, num_classes)
|
||||
targets (torch.Tensor): Ground truth labels of shape (batch_size,)
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Computed focal loss
|
||||
"""
|
||||
# Ensure input types are consistent
|
||||
inputs = inputs.to(torch.float32)
|
||||
targets = targets.to(torch.float32)
|
||||
|
||||
# Handle shape matching
|
||||
if inputs.dim() == 2 and inputs.size(1) == 1:
|
||||
inputs = inputs.squeeze(1)
|
||||
|
||||
# Fully fused computation using highly optimized built-in functions
|
||||
# Compute sigmoid with numerical stability
|
||||
sigmoid_inputs = torch.sigmoid(inputs)
|
||||
|
||||
# Compute pt based on targets
|
||||
pt = torch.where(targets == 1, sigmoid_inputs, 1 - sigmoid_inputs)
|
||||
|
||||
# Compute focal weight
|
||||
focal_weight = self.alpha * torch.pow(1 - pt, self.gamma)
|
||||
|
||||
# Compute binary cross entropy with logits (more numerically stable)
|
||||
bce = F.binary_cross_entropy_with_logits(inputs, targets, reduction='none')
|
||||
|
||||
# Apply focal weight
|
||||
focal_loss = focal_weight * bce
|
||||
|
||||
# Apply reduction
|
||||
if self.reduction == 'mean':
|
||||
return focal_loss.mean()
|
||||
elif self.reduction == 'sum':
|
||||
return focal_loss.sum()
|
||||
else:
|
||||
return focal_loss
|
||||
batch_size = 32
|
||||
num_classes = 1
|
||||
|
||||
def get_inputs():
|
||||
# Generate random logits with explicit float32
|
||||
inputs = torch.randn(batch_size, num_classes, dtype=torch.float32)
|
||||
# Generate random binary targets (0 or 1) with explicit float32
|
||||
targets = torch.randint(0, 2, (batch_size,), dtype=torch.float32)
|
||||
return [inputs, targets]
|
||||
|
||||
def get_init_inputs():
|
||||
return []
|
||||
|
||||
|
||||
|
||||
Your task is to optimize this Focal Loss with Sigmoid implementation by:
|
||||
|
||||
1. **Complete Operator Fusion**: Combine the sigmoid computation and focal loss calculation into a single CUDA kernel to eliminate intermediate tensor storage and redundant computations. The kernel should compute sigmoid, pt, focal weight, and BCE loss in one fused operation.
|
||||
|
||||
2. **Enhanced Numerical Stability**: Implement numerically stable sigmoid computation with conditional branches for positive/negative logits, use more aggressive epsilon bounds (1e-8), and directly compute BCE without intermediate steps to avoid precision loss.
|
||||
|
||||
3. **Memory Access Optimization**: Minimize global memory access by keeping all intermediate computations (sigmoid, pt, focal_weight, bce) in registers, and ensure coalesced memory access patterns for both logits and targets.
|
||||
|
||||
4. **Optimized Thread Configuration**: Use optimal block size of 256 threads with dynamic grid computation based on batch size, and implement efficient kernel launch parameters.
|
||||
|
||||
5. **Type and Shape Consistency**: Ensure all tensors use float32 for consistency, handle shape matching automatically (squeeze dim=1 when needed), and maintain proper device placement.
|
||||
|
||||
The optimized CUDA kernel should:
|
||||
- Take logits and targets as input (both float32)
|
||||
- Compute sigmoid, pt, focal weight, and BCE loss in a single fused kernel
|
||||
- Use optimized sigmoid computation with numerical stability
|
||||
- Apply more aggressive epsilon clamping (1e-8) for better stability
|
||||
- Directly compute BCE without intermediate probability storage
|
||||
- Output the fused focal loss values
|
||||
- Support both 'mean' and 'sum' reduction modes
|
||||
- Use optimized compilation flags (-O3, --use_fast_math)
|
||||
- Achieve significant speedup (1.4-1.6x) over the original PyTorch implementation through fused computation and reduced memory overhead
|
||||
|
||||
Follow the inline CUDA extension syntax example provided in the reference. The kernel should demonstrate performance improvements through complete operator fusion, enhanced numerical stability, and optimized memory access patterns.
|
||||
|
|
@ -0,0 +1,84 @@
|
|||
###########################################################
|
||||
# 性能和精度验证程序
|
||||
###########################################################
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import time
|
||||
from focalloss_sigmoid_torchcode import Model, get_inputs, get_init_inputs
|
||||
from focalloss_sigmoid_cudacode 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, atol=1e-05)
|
||||
max_diff = torch.max(torch.abs(output_torch - output_cuda)).item()
|
||||
mean_diff = torch.mean(torch.abs(output_torch - output_cuda)).item()
|
||||
|
||||
if precision_flag:
|
||||
print(f"✅ 精度对齐：两个模型的输出结果非常接近。")
|
||||
print(f"最大误差: {max_diff:.8f}, 平均误差: {mean_diff:.8f}")
|
||||
else:
|
||||
print(f"❌ 精度不一致！最大误差: {max_diff:.8f}, 平均误差: {mean_diff:.8f}")
|
||||
|
||||
print("\n-------------------- 性能加速比测试 --------------------")
|
||||
num_iterations = 100
|
||||
|
||||
# GPU 预热
|
||||
for _ in range(10):
|
||||
_ = torch_model(*inputs)
|
||||
_ = cuda_model(*inputs)
|
||||
|
||||
# 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 Focal Loss_sigmoid 平均执行时间: {torch_time:.6f} 秒")
|
||||
print(f"自定义 CUDA Focal Loss_sigmoid 平均执行时间: {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