forked from ccf-ai-infra/GPUCodeForces
Merge pull request 'FEAT:dice_bce #110' (#1043) from wut0n/GPUCodeForces:dice_bce into main
This commit is contained in:
commit
c8ff6d7bb0
|
|
@ -0,0 +1,128 @@
|
|||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
dice_bce_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
// Warp级归约版本 - 融合了Dice Loss和BCE Loss
|
||||
__global__ void dice_bce_kernel_warp_reduction(
|
||||
const float* __restrict__ logits,
|
||||
const float* __restrict__ target,
|
||||
float* __restrict__ intersection_sum,
|
||||
float* __restrict__ pred_sum,
|
||||
float* __restrict__ target_sum,
|
||||
float* __restrict__ bce_sum,
|
||||
int size
|
||||
) {
|
||||
int tid = threadIdx.x;
|
||||
int warp_id = (blockIdx.x * blockDim.x + tid) / 32;
|
||||
int lane_id = tid % 32;
|
||||
|
||||
int elements_per_warp = 1024;
|
||||
int warp_start = warp_id * elements_per_warp;
|
||||
int warp_end = min(warp_start + elements_per_warp, size);
|
||||
|
||||
float reg_intersection = 0.0f;
|
||||
float reg_pred_sum = 0.0f;
|
||||
float reg_target_sum = 0.0f;
|
||||
float reg_bce = 0.0f;
|
||||
|
||||
for (int i = warp_start + lane_id; i < warp_end; i += 32) {
|
||||
float l = logits[i];
|
||||
float t = target[i];
|
||||
|
||||
// --- 关键融合点：计算Sigmoid和BCE ---
|
||||
// 使用数值稳定的log-sigmoid公式: log(sigmoid(x)) = -log(1 + exp(-x))
|
||||
float p = 1.0f / (1.0f + expf(-l)); // Sigmoid
|
||||
float log_sigmoid_val = -logf(1.0f + expf(-l)); // log(sigmoid(l))
|
||||
|
||||
// 累加Dice Loss所需项
|
||||
reg_intersection += p * t;
|
||||
reg_pred_sum += p;
|
||||
reg_target_sum += t;
|
||||
|
||||
// 累加BCE Loss所需项
|
||||
// BCE = -[t * log(p) + (1-t) * log(1-p)]
|
||||
reg_bce -= (t * log_sigmoid_val + (1.0f - t) * logf(1.0f - p));
|
||||
}
|
||||
|
||||
// 对所有累加器进行Warp内归约
|
||||
#pragma unroll
|
||||
for (int offset = 16; offset > 0; offset /= 2) {
|
||||
reg_intersection += __shfl_down_sync(0xffffffff, reg_intersection, offset);
|
||||
reg_pred_sum += __shfl_down_sync(0xffffffff, reg_pred_sum, offset);
|
||||
reg_target_sum += __shfl_down_sync(0xffffffff, reg_target_sum, offset);
|
||||
reg_bce += __shfl_down_sync(0xffffffff, reg_bce, offset);
|
||||
}
|
||||
|
||||
if (lane_id == 0) {
|
||||
atomicAdd(intersection_sum, reg_intersection);
|
||||
atomicAdd(pred_sum, reg_pred_sum);
|
||||
atomicAdd(target_sum, reg_target_sum);
|
||||
atomicAdd(bce_sum, reg_bce);
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor dice_bce_cuda(torch::Tensor logits, torch::Tensor target, float dice_weight, float bce_weight) {
|
||||
auto size = logits.numel();
|
||||
|
||||
// 创建中间结果张量
|
||||
auto intersection = torch::zeros(1, logits.options());
|
||||
auto pred_sum_tensor = torch::zeros(1, logits.options());
|
||||
auto target_sum_tensor = torch::zeros(1, logits.options());
|
||||
auto bce_sum_tensor = torch::zeros(1, logits.options());
|
||||
|
||||
const int block_size = 256;
|
||||
int num_warps_needed = (size + 1024 - 1) / 1024;
|
||||
int num_blocks = (num_warps_needed * 32 + block_size - 1) / block_size;
|
||||
|
||||
dice_bce_kernel_warp_reduction<<<num_blocks, block_size>>>(
|
||||
logits.data_ptr<float>(),
|
||||
target.data_ptr<float>(),
|
||||
intersection.data_ptr<float>(),
|
||||
pred_sum_tensor.data_ptr<float>(),
|
||||
target_sum_tensor.data_ptr<float>(),
|
||||
bce_sum_tensor.data_ptr<float>(),
|
||||
size
|
||||
);
|
||||
|
||||
// 在Host端计算最终损失
|
||||
float epsilon = 1e-6f;
|
||||
float dice_score = (2.0f * intersection.item<float>()) / (pred_sum_tensor.item<float>() + target_sum_tensor.item<float>() + epsilon);
|
||||
float loss_dice = 1.0f - dice_score;
|
||||
|
||||
float loss_bce = bce_sum_tensor.item<float>() / size;
|
||||
|
||||
float loss_total = dice_weight * loss_dice + bce_weight * loss_bce;
|
||||
|
||||
return torch::tensor(loss_total, logits.options());
|
||||
}
|
||||
"""
|
||||
|
||||
dice_bce_cpp_source = """
|
||||
torch::Tensor dice_bce_cuda(torch::Tensor logits, torch::Tensor target, float dice_weight, float bce_weight);
|
||||
"""
|
||||
|
||||
dice_bce = load_inline(
|
||||
name="dice_bce",
|
||||
cpp_sources=dice_bce_cpp_source,
|
||||
cuda_sources=dice_bce_source,
|
||||
functions=["dice_bce_cuda"],
|
||||
extra_cuda_cflags=[
|
||||
"-O3",
|
||||
"--use_fast_math",
|
||||
"-gencode=arch=compute_80,code=sm_80"
|
||||
],
|
||||
verbose=True
|
||||
)
|
||||
|
||||
class ModelNew(torch.nn.Module):
|
||||
def __init__(self, dice_weight=1.0, bce_weight=1.0):
|
||||
super(ModelNew, self).__init__()
|
||||
self.dice_weight = dice_weight
|
||||
self.bce_weight = bce_weight
|
||||
self.dice_bce = dice_bce
|
||||
|
||||
def forward(self, logits, target):
|
||||
return self.dice_bce.dice_bce_cuda(logits, target, self.dice_weight, self.bce_weight)
|
||||
|
|
@ -0,0 +1,58 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
PyTorch基准实现：Dice Loss + BCEWithLogitsLoss
|
||||
"""
|
||||
def __init__(self, dice_weight=1.0, bce_weight=1.0):
|
||||
super(Model, self).__init__()
|
||||
self.dice_weight = dice_weight
|
||||
self.bce_weight = bce_weight
|
||||
# BCEWithLogitsLoss内部融合了Sigmoid和BCE，更稳定
|
||||
self.bce_loss_fn = nn.BCEWithLogitsLoss(reduction='mean')
|
||||
|
||||
def forward(self, logits: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Computes the combined Dice and BCE loss.
|
||||
|
||||
Args:
|
||||
logits (torch.Tensor): Raw prediction tensor of shape [N, C, H, W].
|
||||
target (torch.Tensor): Target tensor of same shape as logits.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Combined loss value (scalar).
|
||||
"""
|
||||
# --- 第一步：计算Dice Loss ---
|
||||
pred = torch.sigmoid(logits)
|
||||
pred_flat = pred.view(-1)
|
||||
target_flat = target.view(-1)
|
||||
|
||||
intersection = (pred_flat * target_flat).sum()
|
||||
pred_sum = pred_flat.sum()
|
||||
target_sum = target_flat.sum()
|
||||
|
||||
epsilon = 1e-6
|
||||
dice_score = (2.0 * intersection) / (pred_sum + target_sum + epsilon)
|
||||
loss_dice = 1.0 - dice_score
|
||||
|
||||
# --- 第二步：计算BCEWithLogitsLoss ---
|
||||
# PyTorch的BCEWithLogitsLoss已经非常高效且稳定
|
||||
loss_bce = self.bce_loss_fn(logits, target)
|
||||
|
||||
# --- 第三步：组合损失 ---
|
||||
loss_total = self.dice_weight * loss_dice + self.bce_weight * loss_bce
|
||||
|
||||
return loss_total
|
||||
|
||||
batch_size = 32
|
||||
height, width = 256, 256
|
||||
channels = 1
|
||||
|
||||
def get_inputs():
|
||||
logits = torch.randn(batch_size, channels, height, width)
|
||||
target = torch.randint(0, 2, (batch_size, channels, height, width), dtype=torch.float32)
|
||||
return [logits, target]
|
||||
|
||||
def get_init_inputs():
|
||||
return [] # No special initialization inputs needed
|
||||
|
|
@ -0,0 +1,159 @@
|
|||
You write custom CUDA kernels to replace pytorch operators in given architecture to get speedups. You have complete freedom to choose the set of operators you want to replace. You may make the decision to replace some operators with custom CUDA kernels and leave others unchanged. You may replace multiple operators with custom implementations, consider operator fusion opportunities (combining multiple operators into a single kernel, for example, combining matmul+relu), or algorithmic changes (such as online softmax). You are only limited by your imagination.
|
||||
|
||||
Here's an example to show you the syntax of inline embedding custom CUDA operators in PyTorch. The example given architecture is a simple addition:
|
||||
|
||||
python
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
class Model(nn.Module):
|
||||
def init(self) -> None:
|
||||
super().init()
|
||||
|
||||
def forward(self, a, b):
|
||||
return a + b
|
||||
def get_inputs():
|
||||
a = torch.randn(1, 128).cuda()
|
||||
b = torch.randn(1, 128).cuda()
|
||||
return [a, b]
|
||||
|
||||
def get_init_inputs():
|
||||
return []
|
||||
|
||||
|
||||
|
||||
The example new architecture with a custom CUDA kernel looks like this:
|
||||
|
||||
python
|
||||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
add_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
global void add_kernel(const float* a, const float* b, float* out, int size) {
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx < size) {
|
||||
out[idx] = a[idx] + b[idx];
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor add_cuda(torch::Tensor a, torch::Tensor b) {
|
||||
auto out = torch::empty_like(a);
|
||||
int size = a.numel();
|
||||
const int block_size = 256;
|
||||
int num_blocks = (size + block_size - 1) / block_size;
|
||||
add_kernel<<<num_blocks, block_size>>>(a.data_ptr<float>(), b.data_ptr<float>(), out.data_ptr<float>(), size);
|
||||
return out;
|
||||
}
|
||||
"""
|
||||
|
||||
add_cpp_source = """
|
||||
torch::Tensor add_cuda(torch::Tensor a, torch::Tensor b);
|
||||
"""
|
||||
|
||||
Compile the inline CUDA code
|
||||
add = load_inline(
|
||||
name="add",
|
||||
cpp_sources=add_cpp_source,
|
||||
cuda_sources=add_source,
|
||||
functions=["add_cuda"],
|
||||
verbose=True
|
||||
)
|
||||
|
||||
class ModelNew(torch.nn.Module):
|
||||
def init(self):
|
||||
super(ModelNew, self).init()
|
||||
self.add = add
|
||||
|
||||
def forward(self, a, b):
|
||||
return self.add.add_cuda(a, b)
|
||||
|
||||
|
||||
---
|
||||
|
||||
Now, you are given the following PyTorch architecture to accelerate. The model computes a combined loss, which is a weighted sum of Dice Loss and Binary Cross-Entropy with Logits Loss (BCEWithLogitsLoss). This is a very common pattern in segmentation tasks. The baseline implementation uses separate, highly optimized PyTorch functions for each part.
|
||||
|
||||
python
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
PyTorch基准实现：Dice Loss + BCEWithLogitsLoss
|
||||
"""
|
||||
def __init__(self, dice_weight=1.0, bce_weight=1.0):
|
||||
super(Model, self).__init__()
|
||||
self.dice_weight = dice_weight
|
||||
self.bce_weight = bce_weight
|
||||
# BCEWithLogitsLoss内部融合了Sigmoid和BCE，更稳定
|
||||
self.bce_loss_fn = nn.BCEWithLogitsLoss(reduction='mean')
|
||||
|
||||
def forward(self, logits: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Computes the combined Dice and BCE loss.
|
||||
|
||||
Args:
|
||||
logits (torch.Tensor): Raw prediction tensor of shape [N, C, H, W].
|
||||
target (torch.Tensor): Target tensor of same shape as logits.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Combined loss value (scalar).
|
||||
"""
|
||||
# --- 第一步：计算Dice Loss ---
|
||||
pred = torch.sigmoid(logits)
|
||||
pred_flat = pred.view(-1)
|
||||
target_flat = target.view(-1)
|
||||
|
||||
intersection = (pred_flat * target_flat).sum()
|
||||
pred_sum = pred_flat.sum()
|
||||
target_sum = target_flat.sum()
|
||||
|
||||
epsilon = 1e-6
|
||||
dice_score = (2.0 * intersection) / (pred_sum + target_sum + epsilon)
|
||||
loss_dice = 1.0 - dice_score
|
||||
|
||||
# --- 第二步：计算BCEWithLogitsLoss ---
|
||||
# PyTorch的BCEWithLogitsLoss已经非常高效且稳定
|
||||
loss_bce = self.bce_loss_fn(logits, target)
|
||||
|
||||
# --- 第三步：组合损失 ---
|
||||
loss_total = self.dice_weight * loss_dice + self.bce_weight * loss_bce
|
||||
|
||||
return loss_total
|
||||
|
||||
batch_size = 32
|
||||
height, width = 256, 256
|
||||
channels = 1
|
||||
|
||||
def get_inputs():
|
||||
logits = torch.randn(batch_size, channels, height, width)
|
||||
target = torch.randint(0, 2, (batch_size, channels, height, width), dtype=torch.float32)
|
||||
return [logits, target]
|
||||
|
||||
def get_init_inputs():
|
||||
return [] # No special initialization inputs needed
|
||||
|
||||
|
||||
|
||||
Your task is to generate the `ModelNew` architecture with a single, highly optimized custom CUDA kernel that fuses the computation of both the Dice Loss and the BCEWithLogitsLoss.
|
||||
|
||||
**CRITICAL REQUIREMENTS:**
|
||||
|
||||
1. **Performance Optimization & Fusion:**
|
||||
* **Operator Fusion:** The entire calculation for both losses must be performed within a **single CUDA kernel**. The kernel should take `logits` and `target` as input and compute all necessary intermediate sums.
|
||||
* The kernel should use a **Warp-level reduction** strategy for optimal performance. Warps should collaboratively iterate over the total number of elements.
|
||||
* **Data Reuse:** The Sigmoid of `logits` should be computed once per element and then reused for both the Dice and BCE parts of the calculation. This avoids redundant computation and memory traffic.
|
||||
|
||||
2. **Kernel Logic:**
|
||||
* The kernel should launch a grid of blocks where the total number of threads is sufficient to cover all elements in the input tensors (`logits.numel()`).
|
||||
* Each thread should process one element. Inside the loop, it should compute `p = sigmoid(logits[i])`.
|
||||
* It should then accumulate the necessary values for the final loss calculation: `intersection += p * target[i]`, `pred_sum += p`, `target_sum += target[i]`, and `bce_sum += BCE_loss(p, target[i])`.
|
||||
* Use `__shfl_down_sync` for efficient warp-level reduction to compute the final four sums: `intersection`, `pred_sum`, `target_sum`, and `bce_sum`.
|
||||
* The host-side C++ wrapper should then compute the final `loss_dice`, `loss_bce`, and `loss_total` from these four sums. Use the numerically stable formula for BCE: `-[t * log(p) + (1-t) * log(1-p)]`.
|
||||
|
||||
3. **Code Structure:** Follow the exact structure of the provided example, including `load_inline`, the CUDA source string, the C++ wrapper source string, and the `ModelNew` class. The `get_init_inputs` function must return `[]` to match the baseline.
|
||||
|
||||
4. **Compilation Flags:** Use `-O3` and `--use_fast_math` for maximum performance, as this is a common practice for high-throughput kernels like this. Avoid hardcoding compute capabilities to ensure portability.
|
||||
|
|
@ -0,0 +1,74 @@
|
|||
###########################################################
|
||||
# 性能和精度验证程序
|
||||
###########################################################
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import time
|
||||
from dice_bce_torchcode import Model,get_inputs,get_init_inputs
|
||||
from dice_bce_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)
|
||||
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 dice_bceLoss 平均执行时间: {torch_time:.6f} 秒")
|
||||
print(f"自定义 CUDA dice_bceLoss 平均执行时间: {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