From 15dc90734a8b2990568854c1ce7e650e9ff80fd4 Mon Sep 17 00:00:00 2001 From: hli28146 Date: Mon, 8 Dec 2025 15:29:42 +0800 Subject: [PATCH] finish tverskyloss #57 --- S1/hli28146_#57/prompt.txt | 118 +++++++++++++++++ S1/hli28146_#57/run_code.py | 74 +++++++++++ S1/hli28146_#57/tverskyloss_cuda.py | 184 +++++++++++++++++++++++++++ S1/hli28146_#57/tverskyloss_torch.py | 75 +++++++++++ 4 files changed, 451 insertions(+) create mode 100644 S1/hli28146_#57/prompt.txt create mode 100644 S1/hli28146_#57/run_code.py create mode 100644 S1/hli28146_#57/tverskyloss_cuda.py create mode 100644 S1/hli28146_#57/tverskyloss_torch.py diff --git a/S1/hli28146_#57/prompt.txt b/S1/hli28146_#57/prompt.txt new file mode 100644 index 0000000..87bc048 --- /dev/null +++ b/S1/hli28146_#57/prompt.txt @@ -0,0 +1,118 @@ +Write a custom CUDA kernel to optimize `Tversky Loss`. + +Formula: +Index_c = (TP_c + smooth) / (TP_c + alpha * FP_c + beta * FN_c + smooth) +Loss = 1 - mean(Index_c) +Where: +- TP_c = sum(p_c * g_c) +- FP_c = sum(p_c * (1 - g_c)) +- FN_c = sum((1 - p_c) * g_c) +- p is softmax probability, g is one-hot target. + +Problem Analysis: +1. Memory Usage: Standard implementation materializes large (N, C, Spatial) tensors for Probability and One-Hot targets. For 3D volumetric data, this is extremely expensive. +2. Bandwidth: Calculating sums for TP, FP, FN involves multiple passes over these large tensors. + +Optimization Strategy: Fused Softmax-Accumulation Kernel + +1. Parallelism: One Block per Sample (Batch element). Threads iterate over spatial positions (Grid-Stride Loop). + +2. On-the-Fly Softmax: + For each spatial voxel: + - Read logits for all classes. + - Compute Softmax (Max + SumExp) locally. + - Read target class index. + +3. Fused Accumulation: + Instead of full TP/FP/FN tensors, accumulate sufficient statistics in registers/shared memory: + - `sum_intersection`: sum(p_c) where c == target. + - `sum_p`: sum(p_c) for all c. + - `sum_g`: count(c == target). + + From these: + - TP_c = sum_intersection[c] + - FP_c = sum_p[c] - TP_c + - FN_c = sum_g[c] - TP_c + +4. Reduction & Composition: + - Perform block-level reduction for the per-class statistics. + - Thread 0 computes the Tversky Index and final Loss. + +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 = 64 +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) + +ALPHA = 0.7 +BETA = 0.3 +SMOOTH = 1e-6 +REDUCTION = 'none' + +class TverskyLoss(nn.Module): + """ + Tversky loss function for image segmentation using 3D fully convolutional deep networks + https://arxiv.org/pdf/1706.05721 + """ + def __init__(self, alpha=0.7, beta=0.3, smooth=1e-6, reduction='mean'): + super(TverskyLoss, self).__init__() + self.alpha = alpha + self.beta = beta + self.smooth = smooth + 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) + + probs = F.softmax(logits, dim=1) + + targets_onehot = F.one_hot(targets.long(), num_classes=logits.shape[1]) + targets_onehot = targets_onehot.permute(0, 4, 1, 2, 3).float() + + dims = (2, 3, 4) + + tp = torch.sum(probs * targets_onehot, dim=dims) + fp = torch.sum(probs * (1.0 - targets_onehot), dim=dims) + fn = torch.sum((1.0 - probs) * targets_onehot, dim=dims) + + # Tversky Index + numerator = tp + self.smooth + denominator = tp + self.alpha * fp + self.beta * fn + self.smooth + + score = numerator / denominator + + # Loss = 1 - mean(score) + loss = 1.0 - score.mean(dim=1) + + if self.reduction == 'mean': + return loss.mean() + elif self.reduction == 'sum': + return loss.sum() + return loss + +class Model(nn.Module): + def __init__(self, alpha=0.7, beta=0.3, smooth=1e-6, reduction='none'): + super(Model, self).__init__() + self.loss_fn = TverskyLoss(alpha, beta, smooth, 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 [ALPHA, BETA, SMOOTH, REDUCTION] \ No newline at end of file diff --git a/S1/hli28146_#57/run_code.py b/S1/hli28146_#57/run_code.py new file mode 100644 index 0000000..d1a6f69 --- /dev/null +++ b/S1/hli28146_#57/run_code.py @@ -0,0 +1,74 @@ +########################################################### +# 性能和精度验证程序 +########################################################### +import torch +import torch.nn as nn +import time +from tverskyloss_torch import Model,get_inputs,get_init_inputs +from tverskyloss_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() \ No newline at end of file diff --git a/S1/hli28146_#57/tverskyloss_cuda.py b/S1/hli28146_#57/tverskyloss_cuda.py new file mode 100644 index 0000000..3530641 --- /dev/null +++ b/S1/hli28146_#57/tverskyloss_cuda.py @@ -0,0 +1,184 @@ +import torch +import torch.nn as nn +from torch.utils.cpp_extension import load_inline + +cpp_source = """ +#include +#include + +torch::Tensor tversky_loss_cuda_forward( + const torch::Tensor& logits, + const torch::Tensor& targets, + float alpha, + float beta, + float smooth, + std::string reduction); +""" + +cuda_source = """ +#include +#include +#include +#include + +#define BLOCK_SIZE 256 +#define MAX_CLASSES 32 + +__global__ void tversky_loss_kernel( + float* __restrict__ output, + const float* __restrict__ logits, + const int64_t* __restrict__ targets, + int C, + int spatial_dim, + float alpha, + float beta, + float smooth) +{ + int batch_idx = blockIdx.x; + int tid = threadIdx.x; + + // Registers for per-class accumulation + // We accumulate: + // acc_tp[c]: sum(p * g) -> Intersection + // acc_p[c]: sum(p) -> Prediction Area + // acc_g[c]: sum(g) -> Ground Truth Area + float acc_tp[MAX_CLASSES]; + float acc_p[MAX_CLASSES]; + float acc_g[MAX_CLASSES]; + + #pragma unroll + for (int c = 0; c < MAX_CLASSES; ++c) { + acc_tp[c] = 0.0f; + acc_p[c] = 0.0f; + acc_g[c] = 0.0f; + } + + // Pointers + const float* b_logits = logits + batch_idx * C * spatial_dim; + const int64_t* b_targets = targets + batch_idx * spatial_dim; + + for (int i = tid; i < spatial_dim; i += blockDim.x) { + int64_t target_label = b_targets[i]; + + // Compute Softmax locally 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; + sum_exp += e; + } + float inv_sum = 1.0f / sum_exp; + + // Update accumulators + for (int c = 0; c < C; ++c) { + float p = local_logits[c] * inv_sum; + float g = (c == target_label) ? 1.0f : 0.0f; + + acc_tp[c] += p * g; + acc_p[c] += p; + acc_g[c] += g; + } + } + + __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_tp[c]); + atomicAdd(&s_stats[1 * C + c], acc_p[c]); + atomicAdd(&s_stats[2 * C + c], acc_g[c]); + } + __syncthreads(); + + if (tid == 0) { + float sum_scores = 0.0f; + for (int c = 0; c < C; ++c) { + float tp = s_stats[0 * C + c]; + float sum_p = s_stats[1 * C + c]; + float sum_g = s_stats[2 * C + c]; + + float fp = sum_p - tp; + float fn = sum_g - tp; + + // Tversky = (TP + s) / (TP + a*FP + b*FN + s) + float num = tp + smooth; + float den = tp + alpha * fp + beta * fn + smooth; + + sum_scores += num / den; + } + + output[batch_idx] = 1.0f - (sum_scores / (float)C); + } +} + +torch::Tensor tversky_loss_cuda_forward( + const torch::Tensor& logits, + const torch::Tensor& targets, + float alpha, + float beta, + float smooth, + 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"); + + int batch_size = logits.size(0); + int num_classes = logits.size(1); + + int spatial_dim = 1; + for (int i=2; i>>( + output.data_ptr(), + logits.data_ptr(), + targets.data_ptr(), + num_classes, + spatial_dim, + alpha, + beta, + smooth + ); + + if (reduction == "mean") { + return output.mean(); + } else if (reduction == "sum") { + return output.sum(); + } + + return output; +} +""" + +class ModelNew(nn.Module): + def __init__(self, alpha=0.7, beta=0.3, smooth=1e-6, reduction='none'): + super(ModelNew, self).__init__() + self.alpha = alpha + self.beta = beta + self.smooth = smooth + self.reduction = reduction + self.op = load_inline( + name='tversky_loss_op', + cpp_sources=cpp_source, + cuda_sources=cuda_source, + functions=['tversky_loss_cuda_forward'], + verbose=False, + extra_cuda_cflags=['-O3'] + ) + + def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor: + return self.op.tversky_loss_cuda_forward(logits.contiguous(), targets.contiguous(), + self.alpha, self.beta, self.smooth, self.reduction) \ No newline at end of file diff --git a/S1/hli28146_#57/tverskyloss_torch.py b/S1/hli28146_#57/tverskyloss_torch.py new file mode 100644 index 0000000..85986f8 --- /dev/null +++ b/S1/hli28146_#57/tverskyloss_torch.py @@ -0,0 +1,75 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +BATCH_SIZE = 16 +NUM_CLASSES = 4 +DEPTH = 64 +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) + +ALPHA = 0.7 +BETA = 0.3 +SMOOTH = 1e-6 +REDUCTION = 'none' + +class TverskyLoss(nn.Module): + """ + Tversky loss function for image segmentation using 3D fully convolutional deep networks + https://arxiv.org/pdf/1706.05721 + """ + def __init__(self, alpha=0.7, beta=0.3, smooth=1e-6, reduction='mean'): + super(TverskyLoss, self).__init__() + self.alpha = alpha + self.beta = beta + self.smooth = smooth + 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) + + probs = F.softmax(logits, dim=1) + + targets_onehot = F.one_hot(targets.long(), num_classes=logits.shape[1]) + targets_onehot = targets_onehot.permute(0, 4, 1, 2, 3).float() + + dims = (2, 3, 4) + + tp = torch.sum(probs * targets_onehot, dim=dims) + fp = torch.sum(probs * (1.0 - targets_onehot), dim=dims) + fn = torch.sum((1.0 - probs) * targets_onehot, dim=dims) + + # Tversky Index + numerator = tp + self.smooth + denominator = tp + self.alpha * fp + self.beta * fn + self.smooth + + score = numerator / denominator + + # Loss = 1 - mean(score) + loss = 1.0 - score.mean(dim=1) + + if self.reduction == 'mean': + return loss.mean() + elif self.reduction == 'sum': + return loss.sum() + return loss + +class Model(nn.Module): + def __init__(self, alpha=0.7, beta=0.3, smooth=1e-6, reduction='none'): + super(Model, self).__init__() + self.loss_fn = TverskyLoss(alpha, beta, smooth, 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 [ALPHA, BETA, SMOOTH, REDUCTION] \ No newline at end of file