forked from ccf-ai-infra/GPUCodeForces
finish GeneralizedDiceLoss #55
This commit is contained in:
parent
f876a28ada
commit
a94de17218
|
|
@ -0,0 +1,183 @@
|
|||
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 gdl_loss_cuda_forward(
|
||||
const torch::Tensor& logits,
|
||||
const torch::Tensor& targets,
|
||||
float eps,
|
||||
std::string reduction);
|
||||
"""
|
||||
|
||||
cuda_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <math.h>
|
||||
#include <cfloat>
|
||||
|
||||
#define BLOCK_SIZE 256
|
||||
// For GDL, C is usually small (Brain: 4-5 classes)
|
||||
#define MAX_CLASSES 32
|
||||
|
||||
struct __align__(16) Float4 {
|
||||
float x, y, z, w;
|
||||
};
|
||||
|
||||
// --- Fused GDL Kernel ---
|
||||
// Grid: (B, 1, 1), Block: (256, 1, 1)
|
||||
// Since C is small, we can keep C accumulators in registers/shared mem per thread.
|
||||
__global__ void gdl_loss_kernel(
|
||||
float* __restrict__ output,
|
||||
const float* __restrict__ logits,
|
||||
const int64_t* __restrict__ targets,
|
||||
int C,
|
||||
int spatial_dim,
|
||||
float eps)
|
||||
{
|
||||
int batch_idx = blockIdx.x;
|
||||
int tid = threadIdx.x;
|
||||
|
||||
float acc_r[MAX_CLASSES];
|
||||
float acc_inter[MAX_CLASSES];
|
||||
float acc_union[MAX_CLASSES];
|
||||
|
||||
#pragma unroll
|
||||
for (int c = 0; c < MAX_CLASSES; ++c) {
|
||||
acc_r[c] = 0.0f;
|
||||
acc_inter[c] = 0.0f;
|
||||
acc_union[c] = 0.0f;
|
||||
}
|
||||
|
||||
// Offset for this batch
|
||||
const float* b_logits = logits + batch_idx * C * spatial_dim;
|
||||
const int64_t* b_targets = targets + batch_idx * spatial_dim;
|
||||
|
||||
// Grid-Stride Loop over spatial pixels
|
||||
for (int i = tid; i < spatial_dim; i += blockDim.x) {
|
||||
// Read Target
|
||||
int64_t target_label = b_targets[i];
|
||||
|
||||
// Read Logits & Compute Softmax for this pixel
|
||||
float local_logits[MAX_CLASSES];
|
||||
float max_l = -FLT_MAX;
|
||||
|
||||
for (int c = 0; c < C; ++c) {
|
||||
float val = b_logits[c * spatial_dim + i];
|
||||
local_logits[c] = val;
|
||||
if (val > max_l) max_l = val;
|
||||
}
|
||||
|
||||
float sum_exp = 0.0f;
|
||||
for (int c = 0; c < C; ++c) {
|
||||
float e = __expf(local_logits[c] - max_l);
|
||||
local_logits[c] = e; // Store exp temporarily
|
||||
sum_exp += e;
|
||||
}
|
||||
|
||||
// 3. Accumulate stats
|
||||
float inv_sum = 1.0f / sum_exp;
|
||||
|
||||
for (int c = 0; c < C; ++c) {
|
||||
float p = local_logits[c] * inv_sum; // prob
|
||||
float r = (c == target_label) ? 1.0f : 0.0f; // one-hot
|
||||
|
||||
acc_r[c] += r;
|
||||
acc_inter[c] += r * p;
|
||||
acc_union[c] += r + p;
|
||||
}
|
||||
}
|
||||
|
||||
// Block Reduction
|
||||
__shared__ float s_stats[3 * MAX_CLASSES];
|
||||
if (tid < 3 * C) {
|
||||
s_stats[tid] = 0.0f;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (int c = 0; c < C; ++c) {
|
||||
atomicAdd(&s_stats[0 * C + c], acc_r[c]);
|
||||
atomicAdd(&s_stats[1 * C + c], acc_inter[c]);
|
||||
atomicAdd(&s_stats[2 * C + c], acc_union[c]);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (tid == 0) {
|
||||
float num = 0.0f;
|
||||
float den = 0.0f;
|
||||
|
||||
for (int c = 0; c < C; ++c) {
|
||||
float sum_r = s_stats[0 * C + c];
|
||||
float sum_inter = s_stats[1 * C + c];
|
||||
float sum_union = s_stats[2 * C + c];
|
||||
|
||||
float w = 1.0f / (sum_r * sum_r + eps);
|
||||
|
||||
num += w * sum_inter;
|
||||
den += w * sum_union;
|
||||
}
|
||||
|
||||
float dice = (2.0f * num) / (den + eps);
|
||||
output[batch_idx] = 1.0f - dice;
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor gdl_loss_cuda_forward(
|
||||
const torch::Tensor& logits,
|
||||
const torch::Tensor& targets,
|
||||
float eps,
|
||||
std::string reduction)
|
||||
{
|
||||
TORCH_CHECK(logits.is_cuda() && targets.is_cuda(), "Inputs must be CUDA");
|
||||
TORCH_CHECK(logits.is_contiguous() && targets.is_contiguous(), "Inputs must be contiguous");
|
||||
|
||||
// (B, C, Spatial)
|
||||
int batch_size = logits.size(0);
|
||||
int num_classes = logits.size(1);
|
||||
|
||||
// Total spatial pixels
|
||||
int spatial_dim = 1;
|
||||
for (int i=2; i<logits.dim(); ++i) spatial_dim *= logits.size(i);
|
||||
|
||||
TORCH_CHECK(num_classes <= MAX_CLASSES, "Num classes exceeds kernel limit");
|
||||
|
||||
auto output = torch::empty({batch_size}, logits.options());
|
||||
|
||||
gdl_loss_kernel<<<batch_size, BLOCK_SIZE>>>(
|
||||
output.data_ptr<float>(),
|
||||
logits.data_ptr<float>(),
|
||||
targets.data_ptr<int64_t>(),
|
||||
num_classes,
|
||||
spatial_dim,
|
||||
eps
|
||||
);
|
||||
|
||||
if (reduction == "mean") {
|
||||
return output.mean();
|
||||
} else if (reduction == "sum") {
|
||||
return output.sum();
|
||||
}
|
||||
|
||||
return output;
|
||||
}
|
||||
"""
|
||||
|
||||
class ModelNew(nn.Module):
|
||||
def __init__(self, eps=1e-6, reduction='none'):
|
||||
super(ModelNew, self).__init__()
|
||||
self.eps = eps
|
||||
self.reduction = reduction
|
||||
self.op = load_inline(
|
||||
name='gdl_loss_op',
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=cuda_source,
|
||||
functions=['gdl_loss_cuda_forward'],
|
||||
verbose=False,
|
||||
extra_cuda_cflags=['-O3']
|
||||
)
|
||||
|
||||
def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
|
||||
return self.op.gdl_loss_cuda_forward(logits.contiguous(), targets.contiguous(), self.eps, self.reduction)
|
||||
|
|
@ -0,0 +1,76 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
BATCH_SIZE = 16
|
||||
NUM_CLASSES = 4
|
||||
DEPTH = 32
|
||||
HEIGHT = 128
|
||||
WIDTH = 128
|
||||
SPATIAL_DIM = DEPTH * HEIGHT * WIDTH
|
||||
SHAPE_LOGITS = (BATCH_SIZE, NUM_CLASSES, DEPTH, HEIGHT, WIDTH)
|
||||
SHAPE_TARGET = (BATCH_SIZE, DEPTH, HEIGHT, WIDTH)
|
||||
|
||||
EPS = 1e-6
|
||||
REDUCTION = 'none'
|
||||
|
||||
class GeneralizedDiceLoss(nn.Module):
|
||||
"""
|
||||
Generalized Dice Loss (Sudre et al. MICCAI 2017)
|
||||
https://arxiv.org/pdf/1707.03237
|
||||
"""
|
||||
def __init__(self, eps=1e-6, reduction='mean'):
|
||||
super(GeneralizedDiceLoss, self).__init__()
|
||||
self.eps = eps
|
||||
self.reduction = reduction
|
||||
|
||||
def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
|
||||
# logits: (B, C, D, H, W)
|
||||
# targets: (B, D, H, W) -> indices
|
||||
probs = F.softmax(logits, dim=1)
|
||||
|
||||
# One-hot encoding
|
||||
targets_onehot = F.one_hot(targets.long(), num_classes=logits.shape[1])
|
||||
targets_onehot = targets_onehot.permute(0, 4, 1, 2, 3).float()
|
||||
|
||||
# Compute weights w_l = 1 / (sum(r_ln)^2)
|
||||
sum_r = targets_onehot.sum(dim=(2, 3, 4)) # (B, C)
|
||||
weights = 1.0 / (sum_r * sum_r + self.eps)
|
||||
|
||||
# Compute Intersection & Union
|
||||
# Intersection: r * p
|
||||
intersection = (targets_onehot * probs).sum(dim=(2, 3, 4)) # (B, C)
|
||||
# Union: r + p
|
||||
union = (targets_onehot + probs).sum(dim=(2, 3, 4)) # (B, C)
|
||||
|
||||
# Weighted Sum
|
||||
# Numerator: 2 * sum_l (w_l * inter_l)
|
||||
# Denominator: sum_l (w_l * union_l)
|
||||
|
||||
numerator = 2.0 * (weights * intersection).sum(dim=1)
|
||||
denominator = (weights * union).sum(dim=1)
|
||||
|
||||
dice_score = numerator / (denominator + self.eps)
|
||||
loss = 1.0 - dice_score
|
||||
|
||||
if self.reduction == 'mean':
|
||||
return loss.mean()
|
||||
elif self.reduction == 'sum':
|
||||
return loss.sum()
|
||||
return loss
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, eps=1e-6, reduction='none'):
|
||||
super(Model, self).__init__()
|
||||
self.loss_fn = GeneralizedDiceLoss(eps=eps, reduction=reduction)
|
||||
|
||||
def forward(self, logits, targets):
|
||||
return self.loss_fn(logits, targets)
|
||||
|
||||
def get_inputs():
|
||||
logits = torch.randn(SHAPE_LOGITS, dtype=torch.float32)
|
||||
targets = torch.randint(0, NUM_CLASSES, SHAPE_TARGET, dtype=torch.long)
|
||||
return [logits.contiguous(), targets.contiguous()]
|
||||
|
||||
def get_init_inputs():
|
||||
return [EPS, REDUCTION]
|
||||
|
|
@ -0,0 +1,122 @@
|
|||
Write a custom CUDA kernel to optimize `Generalized Dice Loss` (GDL).
|
||||
|
||||
Formula:
|
||||
GDL = 1 - 2 * (Sum_l w_l * Sum_n (r_ln * p_ln)) / (Sum_l w_l * Sum_n (r_ln + p_ln))
|
||||
Where:
|
||||
- l is class index, n is spatial index (pixels/voxels).
|
||||
- p_ln is softmax probability.
|
||||
- r_ln is one-hot ground truth.
|
||||
- w_l = 1 / (Sum_n r_ln)^2.
|
||||
|
||||
Problem Analysis:
|
||||
1. Memory Bandwidth: Standard implementation involves Softmax, One-hot generation, Summation per class for weights, and then weighted intersection/union sums. This requires multiple passes over the large (N, C, Spatial) tensors.
|
||||
2. Intermediate Storage: Storing probability maps and one-hot targets consumes significant memory.
|
||||
|
||||
Optimization Strategy: Fused Softmax-Reduction Kernel
|
||||
|
||||
1. One-Block-per-Sample: Each block processes one image/volume in the batch.
|
||||
|
||||
2. On-the-fly Calculation:
|
||||
- Compute Softmax probabilities `p` from logits on-the-fly using cached Max/SumExp.
|
||||
- Generate `r` (one-hot) from target indices on-the-fly.
|
||||
|
||||
3. Fused Accumulation:
|
||||
- Iterate over all spatial pixels `n` and classes `l`.
|
||||
- Maintain 3 accumulators per class in Shared Memory/Registers:
|
||||
- `sum_r`: sum(r_ln) [for weights]
|
||||
- `sum_inter`: sum(r_ln * p_ln)
|
||||
- `sum_union`: sum(r_ln + p_ln)
|
||||
|
||||
Wait, `w_l` depends on `sum_r` across the WHOLE spatial dimension. So we must accumulate `sum_r` for all pixels first?
|
||||
Yes, standard GDL weights depend on the GT volume.
|
||||
However, `sum_inter` and `sum_union` also require summation over `n`.
|
||||
We can accumulate `sum_r`, `sum_inter`, `sum_union` simultaneously in one pass over spatial dimensions.
|
||||
|
||||
4. Shared Memory Reduction:
|
||||
- Use atomicAdd or tree reduction in Shared Memory to aggregate these sums for each class across threads.
|
||||
|
||||
5. Final Composition:
|
||||
- Thread 0 reads the aggregated per-class sums.
|
||||
- Computes `w_l = 1 / (sum_r^2 + eps)`.
|
||||
- Computes numerator `2 * sum(w_l * sum_inter)` and denominator `sum(w_l * sum_union)`.
|
||||
- Writes the final loss per sample.
|
||||
|
||||
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 = 16
|
||||
NUM_CLASSES = 4
|
||||
DEPTH = 32
|
||||
HEIGHT = 128
|
||||
WIDTH = 128
|
||||
SPATIAL_DIM = DEPTH * HEIGHT * WIDTH
|
||||
SHAPE_LOGITS = (BATCH_SIZE, NUM_CLASSES, DEPTH, HEIGHT, WIDTH)
|
||||
SHAPE_TARGET = (BATCH_SIZE, DEPTH, HEIGHT, WIDTH)
|
||||
|
||||
EPS = 1e-6
|
||||
REDUCTION = 'none'
|
||||
|
||||
class GeneralizedDiceLoss(nn.Module):
|
||||
"""
|
||||
Generalized Dice Loss (Sudre et al. MICCAI 2017)
|
||||
https://arxiv.org/pdf/1707.03237
|
||||
"""
|
||||
def __init__(self, eps=1e-6, reduction='mean'):
|
||||
super(GeneralizedDiceLoss, self).__init__()
|
||||
self.eps = eps
|
||||
self.reduction = reduction
|
||||
|
||||
def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
|
||||
# logits: (B, C, D, H, W)
|
||||
# targets: (B, D, H, W) -> indices
|
||||
probs = F.softmax(logits, dim=1)
|
||||
|
||||
# One-hot encoding
|
||||
targets_onehot = F.one_hot(targets.long(), num_classes=logits.shape[1])
|
||||
targets_onehot = targets_onehot.permute(0, 4, 1, 2, 3).float()
|
||||
|
||||
# Compute weights w_l = 1 / (sum(r_ln)^2)
|
||||
sum_r = targets_onehot.sum(dim=(2, 3, 4)) # (B, C)
|
||||
weights = 1.0 / (sum_r * sum_r + self.eps)
|
||||
|
||||
# Compute Intersection & Union
|
||||
# Intersection: r * p
|
||||
intersection = (targets_onehot * probs).sum(dim=(2, 3, 4)) # (B, C)
|
||||
# Union: r + p
|
||||
union = (targets_onehot + probs).sum(dim=(2, 3, 4)) # (B, C)
|
||||
|
||||
# Weighted Sum
|
||||
# Numerator: 2 * sum_l (w_l * inter_l)
|
||||
# Denominator: sum_l (w_l * union_l)
|
||||
|
||||
numerator = 2.0 * (weights * intersection).sum(dim=1)
|
||||
denominator = (weights * union).sum(dim=1)
|
||||
|
||||
dice_score = numerator / (denominator + self.eps)
|
||||
loss = 1.0 - dice_score
|
||||
|
||||
if self.reduction == 'mean':
|
||||
return loss.mean()
|
||||
elif self.reduction == 'sum':
|
||||
return loss.sum()
|
||||
return loss
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, eps=1e-6, reduction='none'):
|
||||
super(Model, self).__init__()
|
||||
self.loss_fn = GeneralizedDiceLoss(eps=eps, reduction=reduction)
|
||||
|
||||
def forward(self, logits, targets):
|
||||
return self.loss_fn(logits, targets)
|
||||
|
||||
def get_inputs():
|
||||
logits = torch.randn(SHAPE_LOGITS, dtype=torch.float32)
|
||||
targets = torch.randint(0, NUM_CLASSES, SHAPE_TARGET, dtype=torch.long)
|
||||
return [logits.contiguous(), targets.contiguous()]
|
||||
|
||||
def get_init_inputs():
|
||||
return [EPS, REDUCTION]
|
||||
|
|
@ -0,0 +1,74 @@
|
|||
###########################################################
|
||||
# 性能和精度验证程序
|
||||
###########################################################
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import time
|
||||
from GDL_torch import Model,get_inputs,get_init_inputs
|
||||
from GDL_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