forked from ccf-ai-infra/GPUCodeForces
feat:add high performance focalloss_reduction #40
This commit is contained in:
commit
e3a4f36376
|
|
@ -0,0 +1,132 @@
|
|||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
import os
|
||||
|
||||
focal_loss_fused_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <cmath>
|
||||
|
||||
__global__ void focal_loss_forward_atomic_kernel(
|
||||
const float* logits,
|
||||
const float* targets,
|
||||
int batch_size,
|
||||
float alpha,
|
||||
float gamma,
|
||||
float* final_loss_ptr
|
||||
) {
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
if (idx < batch_size) {
|
||||
float logit = logits[idx];
|
||||
float target = targets[idx];
|
||||
|
||||
float prob;
|
||||
if (logit > 0) {
|
||||
prob = 1.0f / (1.0f + expf(-logit));
|
||||
} else {
|
||||
float exp_logit = expf(logit);
|
||||
prob = exp_logit / (1.0f + exp_logit);
|
||||
}
|
||||
const float eps = 1e-7f;
|
||||
prob = fmaxf(fminf(prob, 1.0f - eps), eps);
|
||||
|
||||
float one_minus_prob = 1.0f - prob;
|
||||
float log_prob = logf(prob);
|
||||
float log_one_minus_prob = logf(one_minus_prob);
|
||||
|
||||
float pt = (target > 0.5f) ? prob : one_minus_prob;
|
||||
float bce = -(target * log_prob + (1.0f - target) * log_one_minus_prob);
|
||||
|
||||
// --- 关键修复：用exp(log(x))重构pow(x, gamma) ---
|
||||
float one_minus_pt = 1.0f - pt;
|
||||
float focal_weight;
|
||||
if (gamma == 2.0f) {
|
||||
// 对于gamma=2.0，直接乘法更快更精确
|
||||
focal_weight = alpha * one_minus_pt * one_minus_pt;
|
||||
} else {
|
||||
// 对于其他gamma，使用exp(log)方法
|
||||
focal_weight = alpha * expf(gamma * logf(one_minus_pt));
|
||||
}
|
||||
|
||||
float my_loss = focal_weight * bce;
|
||||
|
||||
atomicAdd(final_loss_ptr, my_loss);
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor focal_loss_reduction_fused_cuda(
|
||||
torch::Tensor logits,
|
||||
torch::Tensor targets,
|
||||
float alpha,
|
||||
float gamma,
|
||||
std::string reduction
|
||||
) {
|
||||
logits = logits.contiguous().to(torch::kFloat32);
|
||||
targets = targets.contiguous().to(torch::kFloat32);
|
||||
|
||||
auto batch_size = logits.numel();
|
||||
auto final_loss_tensor = torch::zeros({1}, logits.options());
|
||||
|
||||
const int block_size = 256;
|
||||
int num_blocks = (batch_size + block_size - 1) / block_size;
|
||||
|
||||
focal_loss_forward_atomic_kernel<<<num_blocks, block_size>>>(
|
||||
logits.data_ptr<float>(),
|
||||
targets.data_ptr<float>(),
|
||||
batch_size,
|
||||
alpha,
|
||||
gamma,
|
||||
final_loss_tensor.data_ptr<float>()
|
||||
);
|
||||
|
||||
if (reduction == "mean") {
|
||||
final_loss_tensor = final_loss_tensor / batch_size;
|
||||
}
|
||||
|
||||
return final_loss_tensor;
|
||||
}
|
||||
"""
|
||||
|
||||
focal_loss_fused_cpp_source = """
|
||||
torch::Tensor focal_loss_reduction_fused_cuda(
|
||||
torch::Tensor logits,
|
||||
torch::Tensor targets,
|
||||
float alpha,
|
||||
float gamma,
|
||||
std::string reduction
|
||||
);
|
||||
"""
|
||||
|
||||
build_dir = './cuda_build_focal_reduction_final'
|
||||
os.makedirs(build_dir, exist_ok=True)
|
||||
|
||||
focal_loss_module = load_inline(
|
||||
name="focal_loss_reduction_fused_final",
|
||||
cpp_sources=focal_loss_fused_cpp_source,
|
||||
cuda_sources=focal_loss_fused_source,
|
||||
functions=["focal_loss_reduction_fused_cuda"],
|
||||
verbose=True,
|
||||
build_directory=build_dir,
|
||||
extra_cuda_cflags=["-O3"]
|
||||
)
|
||||
|
||||
class ModelNew(torch.nn.Module):
|
||||
def __init__(self, alpha=0.25, gamma=2.0, reduction='mean'):
|
||||
super(ModelNew, self).__init__()
|
||||
self.alpha = alpha
|
||||
self.gamma = gamma
|
||||
self.reduction = reduction
|
||||
|
||||
def forward(self, logits, targets):
|
||||
logits = logits.to(torch.float32)
|
||||
targets = targets.to(torch.float32)
|
||||
|
||||
if logits.dim() == 2 and logits.size(1) == 1:
|
||||
logits = logits.squeeze(1)
|
||||
|
||||
loss = focal_loss_module.focal_loss_reduction_fused_cuda(
|
||||
logits, targets, self.alpha, self.gamma, self.reduction
|
||||
)
|
||||
|
||||
return loss
|
||||
|
|
@ -0,0 +1,51 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
Baseline Focal Loss using PyTorch's most numerically stable API.
|
||||
"""
|
||||
def __init__(self, alpha=0.25, gamma=2.0, reduction='mean'):
|
||||
super(Model, self).__init__()
|
||||
self.alpha = alpha
|
||||
self.gamma = gamma
|
||||
self.reduction = reduction
|
||||
|
||||
def forward(self, inputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
|
||||
# --- 关键修复1：强制转换数据类型 ---
|
||||
# 确保所有计算都在float32下进行，避免类型错误
|
||||
inputs = inputs.to(torch.float32)
|
||||
targets = targets.to(torch.float32)
|
||||
|
||||
# --- 关键修复2：统一输入维度 ---
|
||||
if inputs.dim() == 2 and inputs.size(1) == 1:
|
||||
inputs = inputs.squeeze(1)
|
||||
|
||||
# --- 使用PyTorch官方推荐的稳定API计算BCE ---
|
||||
# 现在inputs和targets的维度和类型都正确了
|
||||
bce = F.binary_cross_entropy_with_logits(inputs, targets, reduction='none')
|
||||
|
||||
# 后续的Focal Loss计算保持不变
|
||||
probs = torch.sigmoid(inputs)
|
||||
pt = torch.where(targets == 1, probs, 1 - probs)
|
||||
focal_weight = self.alpha * torch.pow(1 - pt, self.gamma)
|
||||
focal_loss = focal_weight * bce
|
||||
|
||||
if self.reduction == 'mean':
|
||||
return focal_loss.mean()
|
||||
elif self.reduction == 'sum':
|
||||
return focal_loss.sum()
|
||||
else:
|
||||
return focal_loss
|
||||
|
||||
batch_size = 1024
|
||||
num_classes = 1
|
||||
|
||||
def get_inputs():
|
||||
inputs = torch.randn(batch_size, num_classes)
|
||||
targets = torch.randint(0, 2, (batch_size,))
|
||||
return [inputs, targets]
|
||||
|
||||
def get_init_inputs():
|
||||
return []
|
||||
|
|
@ -0,0 +1,104 @@
|
|||
You write custom CUDA kernels to replace the pytorch operators in the 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 torch: The example given architecture is:
|
||||
|
||||
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():
|
||||
# randomly generate input tensors based on the model architecture
|
||||
a = torch.randn(1, 128).cuda()
|
||||
b = torch.randn(1, 128).cuda()
|
||||
return [a, b]
|
||||
|
||||
def get_init_inputs():
|
||||
# randomly generate tensors required for initialization based on the model architecture
|
||||
return []
|
||||
|
||||
|
||||
The example new arch with custom CUDA kernels looks like this:
|
||||
|
||||
|
||||
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():
|
||||
# randomly generate input tensors based on the model architecture
|
||||
a = torch.randn(1, 128).cuda()
|
||||
b = torch.randn(1, 128).cuda()
|
||||
return [a, b]
|
||||
|
||||
|
||||
def get_init_inputs():
|
||||
# randomly generate tensors required for initialization based on the model architecture
|
||||
return []
|
||||
You are given the following architecture:
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
Baseline Focal Loss implementation using fundamental PyTorch ops.
|
||||
"""
|
||||
def __init__(self, alpha=0.25, gamma=2.0, reduction='mean'):
|
||||
super(Model, self).__init__()
|
||||
self.alpha = alpha
|
||||
self.gamma = gamma
|
||||
self.reduction = reduction
|
||||
|
||||
def forward(self, inputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
|
||||
# 1. Manually implement Sigmoid
|
||||
probs = torch.sigmoid(inputs)
|
||||
|
||||
# 2. Manually implement Binary Cross Entropy
|
||||
eps = 1e-7
|
||||
bce = -(targets * torch.log(probs + eps) + (1 - targets) * torch.log(1 - probs + eps))
|
||||
|
||||
# 3. Compute pt
|
||||
pt = torch.where(targets == 1, probs, 1 - probs)
|
||||
|
||||
# 4. Compute Focal Weight
|
||||
focal_weight = self.alpha * torch.pow(1 - pt, self.gamma)
|
||||
|
||||
# 5. Apply Focal Weight
|
||||
focal_loss = focal_weight * bce
|
||||
|
||||
if self.reduction == 'mean':
|
||||
return focal_loss.mean()
|
||||
elif self.reduction == 'sum':
|
||||
return focal_loss.sum()
|
||||
else:
|
||||
return focal_loss
|
||||
|
||||
batch_size = 1024
|
||||
num_classes = 1
|
||||
|
||||
def get_inputs():
|
||||
inputs = torch.randn(batch_size, num_classes)
|
||||
targets = torch.randint(0, 2, (batch_size,))
|
||||
return [inputs, targets]
|
||||
|
||||
def get_init_inputs():
|
||||
return []
|
||||
|
|
@ -0,0 +1,74 @@
|
|||
###########################################################
|
||||
# 性能和精度验证程序
|
||||
###########################################################
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import time
|
||||
from focalloss_reduction_torchcode import Model,get_inputs,get_init_inputs
|
||||
from focalloss_reduction_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 focalloss_reduction 平均执行时间: {torch_time:.6f} 秒")
|
||||
print(f"自定义 CUDA focalloss_reduction 平均执行时间: {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