forked from ccf-ai-infra/GPUCodeForces
Merge pull request 'finish tverskyloss #57' (#675) from hli28146/GPUCodeForces:h57 into main
This commit is contained in:
commit
ae6d9721ef
|
|
@ -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]
|
||||
|
|
@ -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()
|
||||
|
|
@ -0,0 +1,184 @@
|
|||
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 tversky_loss_cuda_forward(
|
||||
const torch::Tensor& logits,
|
||||
const torch::Tensor& targets,
|
||||
float alpha,
|
||||
float beta,
|
||||
float smooth,
|
||||
std::string reduction);
|
||||
"""
|
||||
|
||||
cuda_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <math.h>
|
||||
#include <cfloat>
|
||||
|
||||
#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<logits.dim(); ++i) spatial_dim *= logits.size(i);
|
||||
|
||||
TORCH_CHECK(num_classes <= MAX_CLASSES, "Num classes exceeds kernel limit (32)");
|
||||
|
||||
auto output = torch::empty({batch_size}, logits.options());
|
||||
|
||||
tversky_loss_kernel<<<batch_size, BLOCK_SIZE>>>(
|
||||
output.data_ptr<float>(),
|
||||
logits.data_ptr<float>(),
|
||||
targets.data_ptr<int64_t>(),
|
||||
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)
|
||||
|
|
@ -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]
|
||||
Loading…
Reference in New Issue