forked from ccf-ai-infra/GPUCodeForces
Merge pull request 'finish FocalTverskyLoss #58' (#685) from hli28146/GPUCodeForces:h58 into main
This commit is contained in:
commit
f527d82afb
|
|
@ -0,0 +1,208 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
cpp_source = """
|
||||
#include <torch/extension.h>
|
||||
|
||||
torch::Tensor focal_tversky_loss_cuda_forward(
|
||||
const torch::Tensor& y_pred,
|
||||
const torch::Tensor& y_true,
|
||||
float alpha,
|
||||
float beta,
|
||||
float gamma,
|
||||
float epsilon);
|
||||
"""
|
||||
|
||||
cuda_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <math.h>
|
||||
|
||||
#define BLOCK_SIZE 256
|
||||
#define WARP_SIZE 32
|
||||
|
||||
struct __align__(16) Float4 {
|
||||
float x, y, z, w;
|
||||
};
|
||||
|
||||
template<typename T>
|
||||
__device__ __forceinline__ T warp_reduce_sum(T val) {
|
||||
#pragma unroll
|
||||
for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
|
||||
val += __shfl_down_sync(0xffffffff, val, offset);
|
||||
}
|
||||
return val;
|
||||
}
|
||||
|
||||
// Reduce 3 values simultaneously
|
||||
__device__ __forceinline__ void block_reduce_sum_3(
|
||||
float& val1, float& val2, float& val3)
|
||||
{
|
||||
static __shared__ float shared1[32];
|
||||
static __shared__ float shared2[32];
|
||||
static __shared__ float shared3[32];
|
||||
|
||||
int lane = threadIdx.x % WARP_SIZE;
|
||||
int wid = threadIdx.x / WARP_SIZE;
|
||||
|
||||
val1 = warp_reduce_sum(val1);
|
||||
val2 = warp_reduce_sum(val2);
|
||||
val3 = warp_reduce_sum(val3);
|
||||
|
||||
if (lane == 0) {
|
||||
shared1[wid] = val1;
|
||||
shared2[wid] = val2;
|
||||
shared3[wid] = val3;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Last warp reduction (assuming block size 256)
|
||||
val1 = (threadIdx.x < blockDim.x / WARP_SIZE) ? shared1[lane] : 0.0f;
|
||||
val2 = (threadIdx.x < blockDim.x / WARP_SIZE) ? shared2[lane] : 0.0f;
|
||||
val3 = (threadIdx.x < blockDim.x / WARP_SIZE) ? shared3[lane] : 0.0f;
|
||||
|
||||
if (wid == 0) {
|
||||
val1 = warp_reduce_sum(val1);
|
||||
val2 = warp_reduce_sum(val2);
|
||||
val3 = warp_reduce_sum(val3);
|
||||
}
|
||||
}
|
||||
|
||||
// Sigmoid
|
||||
__device__ __forceinline__ float fast_sigmoid(float x) {
|
||||
return 1.0f / (1.0f + __expf(-x));
|
||||
}
|
||||
|
||||
// Fused Kernel
|
||||
__global__ void focal_tversky_kernel(
|
||||
float* __restrict__ output,
|
||||
const float* __restrict__ logits,
|
||||
const float* __restrict__ targets,
|
||||
int spatial_size,
|
||||
int num_blocks,
|
||||
float alpha,
|
||||
float beta,
|
||||
float gamma,
|
||||
float epsilon)
|
||||
{
|
||||
int slice_idx = blockIdx.x;
|
||||
if (slice_idx >= num_blocks) return;
|
||||
|
||||
int offset = slice_idx * spatial_size;
|
||||
const float* slice_logits = logits + offset;
|
||||
const float* slice_targets = targets + offset;
|
||||
|
||||
float tp_sum = 0.0f;
|
||||
float fn_sum = 0.0f;
|
||||
float fp_sum = 0.0f;
|
||||
|
||||
int tid = threadIdx.x;
|
||||
int stride = blockDim.x;
|
||||
|
||||
int i = tid * 4;
|
||||
while (i < spatial_size) {
|
||||
if (i + 4 <= spatial_size) {
|
||||
Float4 l_vec = reinterpret_cast<const Float4*>(&slice_logits[i])[0];
|
||||
Float4 t_vec = reinterpret_cast<const Float4*>(&slice_targets[i])[0];
|
||||
|
||||
// Unroll manually
|
||||
float p, g;
|
||||
|
||||
p = fast_sigmoid(l_vec.x); g = t_vec.x;
|
||||
tp_sum += p * g; fn_sum += g * (1.0f - p); fp_sum += p * (1.0f - g);
|
||||
|
||||
p = fast_sigmoid(l_vec.y); g = t_vec.y;
|
||||
tp_sum += p * g; fn_sum += g * (1.0f - p); fp_sum += p * (1.0f - g);
|
||||
|
||||
p = fast_sigmoid(l_vec.z); g = t_vec.z;
|
||||
tp_sum += p * g; fn_sum += g * (1.0f - p); fp_sum += p * (1.0f - g);
|
||||
|
||||
p = fast_sigmoid(l_vec.w); g = t_vec.w;
|
||||
tp_sum += p * g; fn_sum += g * (1.0f - p); fp_sum += p * (1.0f - g);
|
||||
|
||||
} else {
|
||||
for (int k = 0; k < 4 && i + k < spatial_size; ++k) {
|
||||
float p = fast_sigmoid(slice_logits[i+k]);
|
||||
float g = slice_targets[i+k];
|
||||
tp_sum += p * g;
|
||||
fn_sum += g * (1.0f - p);
|
||||
fp_sum += (1.0f - g) * p;
|
||||
}
|
||||
}
|
||||
i += stride * 4;
|
||||
}
|
||||
|
||||
block_reduce_sum_3(tp_sum, fn_sum, fp_sum);
|
||||
|
||||
if (tid == 0) {
|
||||
float denominator = tp_sum + alpha * fn_sum + beta * fp_sum + epsilon;
|
||||
float ti = (tp_sum + epsilon) / denominator;
|
||||
|
||||
float loss = __powf(1.0f - ti, gamma);
|
||||
output[slice_idx] = loss;
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor focal_tversky_loss_cuda_forward(
|
||||
const torch::Tensor& logits,
|
||||
const torch::Tensor& targets,
|
||||
float alpha,
|
||||
float beta,
|
||||
float gamma,
|
||||
float epsilon)
|
||||
{
|
||||
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 channels = logits.size(1);
|
||||
|
||||
int64_t total_elements = logits.numel();
|
||||
int64_t num_slices = batch_size * channels;
|
||||
int spatial_size = total_elements / num_slices;
|
||||
|
||||
auto output = torch::empty({num_slices}, logits.options());
|
||||
|
||||
focal_tversky_kernel<<<num_slices, BLOCK_SIZE>>>(
|
||||
output.data_ptr<float>(),
|
||||
logits.data_ptr<float>(),
|
||||
targets.data_ptr<float>(),
|
||||
spatial_size,
|
||||
num_slices,
|
||||
alpha,
|
||||
beta,
|
||||
gamma,
|
||||
epsilon
|
||||
);
|
||||
|
||||
return output.mean();
|
||||
}
|
||||
"""
|
||||
|
||||
class ModelNew(nn.Module):
|
||||
def __init__(self, alpha=0.7, beta=0.3, gamma=0.75, epsilon=1e-7):
|
||||
super(ModelNew, self).__init__()
|
||||
self.alpha = alpha
|
||||
self.beta = beta
|
||||
self.gamma = gamma
|
||||
self.epsilon = epsilon
|
||||
|
||||
self.op = load_inline(
|
||||
name='focal_tversky_op',
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=cuda_source,
|
||||
functions=['focal_tversky_loss_cuda_forward'],
|
||||
verbose=False,
|
||||
extra_cuda_cflags=['-O3']
|
||||
)
|
||||
|
||||
def forward(self, y_pred, y_true):
|
||||
return self.op.focal_tversky_loss_cuda_forward(
|
||||
y_pred.contiguous(),
|
||||
y_true.contiguous(),
|
||||
self.alpha,
|
||||
self.beta,
|
||||
self.gamma,
|
||||
self.epsilon
|
||||
)
|
||||
|
|
@ -0,0 +1,62 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
BATCH_SIZE = 16
|
||||
CHANNELS = 4
|
||||
DEPTH = 32
|
||||
HEIGHT = 128
|
||||
WIDTH = 128
|
||||
SHAPE = (BATCH_SIZE, CHANNELS, DEPTH, HEIGHT, WIDTH)
|
||||
|
||||
ALPHA = 0.7
|
||||
BETA = 0.3
|
||||
GAMMA = 0.75
|
||||
EPSILON = 1e-7
|
||||
|
||||
class FocalTverskyLoss(nn.Module):
|
||||
'''
|
||||
FOCAL TVERSKY LOSS(https://arxiv.org/pdf/1810.07842)
|
||||
The input tensors are expected to have a shape of (B, N, H, W, L), where:
|
||||
B is the batch size
|
||||
N is the number of channels
|
||||
H, W, L represent the depth, height, and width of the volumes, respectively
|
||||
'''
|
||||
def __init__(self, alpha=0.7, beta=0.3, gamma=0.75, epsilon=1e-7):
|
||||
super(FocalTverskyLoss, self).__init__()
|
||||
self.alpha = alpha
|
||||
self.beta = beta
|
||||
self.gamma = gamma
|
||||
self.epsilon = epsilon
|
||||
|
||||
def forward(self, y_pred, y_true):
|
||||
# y_pred: logits
|
||||
# y_true: binary targets (0 or 1)
|
||||
y_pred = torch.sigmoid(y_pred)
|
||||
|
||||
# Reduction over spatial dims (2, 3, 4)
|
||||
tp = (y_true * y_pred).sum(dim=(2, 3, 4))
|
||||
fn = (y_true * (1 - y_pred)).sum(dim=(2, 3, 4))
|
||||
fp = ((1 - y_true) * y_pred).sum(dim=(2, 3, 4))
|
||||
|
||||
tversky_index = (tp + self.epsilon) / (tp + self.alpha * fn + self.beta * fp + self.epsilon)
|
||||
|
||||
loss = (1 - tversky_index).pow(self.gamma)
|
||||
|
||||
return loss.mean()
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, alpha=0.7, beta=0.3, gamma=0.75, epsilon=1e-7):
|
||||
super(Model, self).__init__()
|
||||
self.loss_fn = FocalTverskyLoss(alpha, beta, gamma, epsilon)
|
||||
|
||||
def forward(self, y_pred, y_true):
|
||||
return self.loss_fn(y_pred, y_true)
|
||||
|
||||
def get_inputs():
|
||||
y_pred = torch.randn(SHAPE, dtype=torch.float32)
|
||||
y_true = torch.randint(0, 2, SHAPE, dtype=torch.float32)
|
||||
return [y_pred.contiguous(), y_true.contiguous()]
|
||||
|
||||
def get_init_inputs():
|
||||
return [ALPHA, BETA, GAMMA, EPSILON]
|
||||
|
|
@ -0,0 +1,92 @@
|
|||
Write a custom CUDA kernel to optimize `Focal Tversky Loss` for 3D segmentation.
|
||||
|
||||
Formula:
|
||||
TI_c = (TP_c + epsilon) / (TP_c + alpha * FN_c + beta * FP_c + epsilon)
|
||||
Loss = Mean( (1 - TI_c)^gamma )
|
||||
|
||||
Parameters:
|
||||
- alpha, beta, gamma: Loss weighting parameters.
|
||||
- epsilon: Smoothing factor for numerical stability.
|
||||
|
||||
Optimization Strategy: Fused Block-per-Channel Reduction
|
||||
|
||||
1. Flattened View: Treat the input (B, C, D, H, W) as `B * C` independent slices.
|
||||
|
||||
2. Block-per-Slice Parallelism: Launch `B * C` CUDA blocks. Each block reduces one spatial slice to compute TP, FN, and FP simultaneously.
|
||||
|
||||
3. Fused Accumulation:
|
||||
- Load logit and target using vectorized `float4`.
|
||||
- Compute `p = sigmoid(logit)` in register.
|
||||
- Accumulate TP, FN, FP in registers.
|
||||
|
||||
4. Shared Memory Reduction: Perform parallel reduction for the three accumulators.
|
||||
|
||||
5. Final Calculation: Thread 0 uses the reduced TP, FN, FP and the passed `epsilon` to calculate the loss for the slice.
|
||||
|
||||
6. Global Reduction: The C++ wrapper performs the final mean reduction.
|
||||
|
||||
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
|
||||
CHANNELS = 4
|
||||
DEPTH = 32
|
||||
HEIGHT = 128
|
||||
WIDTH = 128
|
||||
SHAPE = (BATCH_SIZE, CHANNELS, DEPTH, HEIGHT, WIDTH)
|
||||
|
||||
ALPHA = 0.7
|
||||
BETA = 0.3
|
||||
GAMMA = 0.75
|
||||
EPSILON = 1e-7
|
||||
|
||||
class FocalTverskyLoss(nn.Module):
|
||||
'''
|
||||
FOCAL TVERSKY LOSS(https://arxiv.org/pdf/1810.07842)
|
||||
The input tensors are expected to have a shape of (B, N, H, W, L), where:
|
||||
B is the batch size
|
||||
N is the number of channels
|
||||
H, W, L represent the depth, height, and width of the volumes, respectively
|
||||
'''
|
||||
def __init__(self, alpha=0.7, beta=0.3, gamma=0.75, epsilon=1e-7):
|
||||
super(FocalTverskyLoss, self).__init__()
|
||||
self.alpha = alpha
|
||||
self.beta = beta
|
||||
self.gamma = gamma
|
||||
self.epsilon = epsilon
|
||||
|
||||
def forward(self, y_pred, y_true):
|
||||
# y_pred: logits
|
||||
# y_true: binary targets (0 or 1)
|
||||
y_pred = torch.sigmoid(y_pred)
|
||||
|
||||
# Reduction over spatial dims (2, 3, 4)
|
||||
tp = (y_true * y_pred).sum(dim=(2, 3, 4))
|
||||
fn = (y_true * (1 - y_pred)).sum(dim=(2, 3, 4))
|
||||
fp = ((1 - y_true) * y_pred).sum(dim=(2, 3, 4))
|
||||
|
||||
tversky_index = (tp + self.epsilon) / (tp + self.alpha * fn + self.beta * fp + self.epsilon)
|
||||
|
||||
loss = (1 - tversky_index).pow(self.gamma)
|
||||
|
||||
return loss.mean()
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, alpha=0.7, beta=0.3, gamma=0.75, epsilon=1e-7):
|
||||
super(Model, self).__init__()
|
||||
self.loss_fn = FocalTverskyLoss(alpha, beta, gamma, epsilon)
|
||||
|
||||
def forward(self, y_pred, y_true):
|
||||
return self.loss_fn(y_pred, y_true)
|
||||
|
||||
def get_inputs():
|
||||
y_pred = torch.randn(SHAPE, dtype=torch.float32)
|
||||
y_true = torch.randint(0, 2, SHAPE, dtype=torch.float32)
|
||||
return [y_pred.contiguous(), y_true.contiguous()]
|
||||
|
||||
def get_init_inputs():
|
||||
return [ALPHA, BETA, GAMMA, EPSILON]
|
||||
|
|
@ -0,0 +1,74 @@
|
|||
###########################################################
|
||||
# 性能和精度验证程序
|
||||
###########################################################
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import time
|
||||
from focaltverskyloss_torch import Model,get_inputs,get_init_inputs
|
||||
from focaltverskyloss_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