feat:add focal loss optimization

S1/wut0n_#2/focalloss_cudacode.py

@@ -0,0 +1,126 @@
+import torch
+from torch.utils.cpp_extension import load_inline
+import os
+
+focal_loss_source = """
+#include <torch/extension.h>
+#include <cuda_runtime.h>
+#include <cmath>
+
+__global__ void focal_loss_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];  // 改为float类型
+        
+        // Compute probability using sigmoid with numerical stability
+        float prob;
+        if (logit > 0) {
+            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)
+        prob = fmaxf(fminf(prob, 1.0f - 1e-7f), 1e-7f);
+        
+        // Compute pt
+        float pt = (target > 0.5f) ? prob : (1.0f - prob);  // 改为float比较
+        
+        // Compute binary cross entropy
+        float bce = -(target * logf(prob) + (1.0f - target) * logf(1.0f - prob));
+        
+        // Compute focal weight
+        float focal_weight = alpha * powf(1.0f - pt, gamma);
+        
+        // Apply focal weight
+        loss[idx] = focal_weight * bce;
+    }
+}
+
+torch::Tensor focal_loss_cuda(torch::Tensor logits, torch::Tensor targets, float alpha, float gamma) {
+    // 确保输入输出数据类型完全一致
+    logits = logits.to(torch::kFloat32);
+    targets = targets.to(torch::kFloat32);  // 改为float32
+    
+    auto batch_size = logits.numel();
+    auto loss = torch::empty_like(logits);
+    
+    // 动态计算线程块数量
+    const int block_size = 256;
+    int num_blocks = (batch_size + block_size - 1) / block_size;
+    
+    // 启动核函数
+    focal_loss_kernel<<<num_blocks, block_size>>>(
+        logits.data_ptr<float>(),
+        targets.data_ptr<float>(),  // 改为float*
+        loss.data_ptr<float>(),
+        batch_size,
+        alpha,
+        gamma
+    );
+    
+    // 检查CUDA错误
+    cudaError_t err = cudaGetLastError();
+    if (err != cudaSuccess) {
+        AT_ERROR("CUDA kernel failed: ", cudaGetErrorString(err));
+    }
+    
+    return loss;
+}
+"""
+
+focal_loss_cpp_source = """
+torch::Tensor focal_loss_cuda(torch::Tensor logits, torch::Tensor targets, float alpha, float gamma);
+"""
+
+
+build_dir = './cuda_build'
+os.makedirs(build_dir, exist_ok=True)
+
+# Compile the inline CUDA code
+focal_loss_module = load_inline(
+    name="focal_loss",
+    cpp_sources=focal_loss_cpp_source,
+    cuda_sources=focal_loss_source,
+    functions=["focal_loss_cuda"],
+    verbose=True,
+    build_directory=build_dir,
+    extra_cuda_cflags=["-O3"] 
+)
+
+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 = focal_loss_module
+
+    def forward(self, logits, targets):
+        # 确保输入输出类型完全一致
+        logits = logits.to(torch.float32)
+        targets = targets.to(torch.float32)  # 改为float32
+        
+        # 确保logits和targets形状一致
+        if logits.dim() == 2 and logits.size(1) == 1:
+            logits = logits.squeeze(1)  # 将logits从[batch_size, 1]变为[batch_size]
+        
+        # Compute focal loss using custom CUDA kernel
+        loss = self.focal_loss.focal_loss_cuda(logits, targets, self.alpha, self.gamma)
+        
+        if self.reduction == 'mean':
+            return loss.mean()
+        elif self.reduction == 'sum':
+            return loss.sum()
+        else:
+            return loss

S1/wut0n_#2/focalloss_torchcode.py

@@ -0,0 +1,69 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class Model(nn.Module):
+    """
+    Focal Loss implementation for binary classification.
+    Focal Loss = -α * (1-pt)^γ * log(pt)
+    where pt = p if target=1, else (1-p)
+    """
+    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 between inputs and targets.
+        
+        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
+        """
+        # 确保输入输出类型完全一致
+        inputs = inputs.to(torch.float32)
+        targets = targets.to(torch.float32)
+        
+        # 确保targets的形状与inputs匹配
+        if inputs.dim() == 2 and inputs.size(1) == 1:
+            targets = targets.unsqueeze(1)  # 将targets从[batch_size]变为[batch_size, 1]
+        
+        # Convert logits to probabilities
+        probs = torch.sigmoid(inputs)
+        
+        # Compute pt
+        pt = torch.where(targets == 1, probs, 1 - probs)
+        
+        # Compute focal weight
+        focal_weight = self.alpha * torch.pow(1 - pt, self.gamma)
+        
+        # Compute binary cross entropy
+        bce = F.binary_cross_entropy_with_logits(inputs, targets, reduction='none')
+        
+        # Apply focal weight
+        focal_loss = focal_weight * bce
+        
+        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 - 明确指定float32
+    inputs = torch.randn(batch_size, num_classes, dtype=torch.float32)
+    # Generate random binary targets (0 or 1) - 明确指定float32
+    targets = torch.randint(0, 2, (batch_size,), dtype=torch.float32)
+    return [inputs, targets]
+
+def get_init_inputs():
+    return []  # No special initialization inputs needed

S1/wut0n_#2/prompt.txt

@@ -0,0 +1,88 @@
+Write a custom CUDA kernel to replace PyTorch's Focal Loss 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 implementation for binary classification.
+Focal Loss = -α * (1-pt)^γ * log(pt)
+where pt = p if target=1, else (1-p)
+“”"
+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 between inputs and targets.
+    
+    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
+    """
+    # Convert logits to probabilities
+    probs = torch.sigmoid(inputs)
+    
+    # Compute pt
+    pt = torch.where(targets == 1, probs, 1 - probs)
+    
+    # Compute focal weight
+    focal_weight = self.alpha * torch.pow(1 - pt, self.gamma)
+    
+    # Compute binary cross entropy
+    bce = F.binary_cross_entropy_with_logits(inputs, targets.float(), reduction='none')
+    
+    # Apply focal weight
+    focal_loss = focal_weight * bce
+    
+    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():
+inputs = torch.randn(batch_size, num_classes, dtype=torch.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 implementation by:
+
+1. **Operator Fusion**: Combine the multiple PyTorch operations (sigmoid, where, pow, binary_cross_entropy_with_logits, multiplication) into a single CUDA kernel to eliminate intermediate tensor storage and memory bandwidth overhead.
+
+2. **Numerical Stability**: Implement numerically stable sigmoid computation to avoid overflow/underflow issues, and add proper bounds checking to prevent log(0) errors.
+
+3. **Memory Access Optimization**: Minimize global memory access by keeping intermediate computations in registers, and ensure coalesced memory access patterns.
+
+4. **Thread Configuration**: Use optimal block size (e.g., 256 threads) and compute grid dimensions dynamically based on input size.
+
+5. **Type Consistency**: Ensure all tensors use float32 for consistency and performance.
+
+The optimized CUDA kernel should:
+- Take logits and targets as input (both float32)
+- Compute sigmoid, pt, focal weight, and BCE loss in a single kernel
+- Output the focal loss values
+- Support both 'mean' and 'sum' reduction modes
+- Maintain numerical stability with proper epsilon handling
+- Achieve significant speedup over the PyTorch implementation
+
+Follow the inline CUDA extension syntax example provided in the reference. The kernel should be optimized for GPU architectures and demonstrate performance improvements through reduced memory access and fused computation.

S1/wut0n_#2/run_code.py

@@ -0,0 +1,84 @@
+###########################################################
+# 性能和精度验证程序
+###########################################################
+import torch
+import torch.nn as nn
+import time
+from focalloss_torchcode import Model, get_inputs, get_init_inputs
+from focalloss_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 平均执行时间: {torch_time:.6f} 秒")
+    print(f"自定义 CUDA Focal Loss 平均执行时间: {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()