forked from ccf-ai-infra/GPUCodeForces
finish centerNetLoss #118
This commit is contained in:
parent
10eed82956
commit
b947ead36d
|
|
@ -0,0 +1,198 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
cuda_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <math.h>
|
||||
|
||||
__inline__ __device__ float warp_reduce(float val) {
|
||||
for (int offset = 16; offset > 0; offset /= 2)
|
||||
val += __shfl_down_sync(0xffffffff, val, offset);
|
||||
return val;
|
||||
}
|
||||
|
||||
__global__ void centernet_hm_loss_kernel(
|
||||
const float* __restrict__ pred_hm,
|
||||
const float* __restrict__ gt_hm,
|
||||
float* __restrict__ global_buffer,
|
||||
int n_elements)
|
||||
{
|
||||
int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int stride = gridDim.x * blockDim.x;
|
||||
|
||||
float local_pos = 0.0f;
|
||||
float local_neg = 0.0f;
|
||||
float local_num = 0.0f;
|
||||
|
||||
for (int idx = tid; idx < n_elements; idx += stride) {
|
||||
float p = pred_hm[idx];
|
||||
float t = gt_hm[idx];
|
||||
|
||||
if (p < 1e-6f) p = 1e-6f;
|
||||
if (p > 0.999999f) p = 0.999999f;
|
||||
|
||||
if (t == 1.0f) {
|
||||
float term = (1.0f - p);
|
||||
local_pos += logf(p) * term * term;
|
||||
local_num += 1.0f;
|
||||
} else {
|
||||
float w = (1.0f - t);
|
||||
w = w * w * w * w;
|
||||
local_neg += logf(1.0f - p) * p * p * w;
|
||||
}
|
||||
}
|
||||
|
||||
local_pos = warp_reduce(local_pos);
|
||||
local_neg = warp_reduce(local_neg);
|
||||
local_num = warp_reduce(local_num);
|
||||
|
||||
if ((threadIdx.x % 32) == 0) {
|
||||
atomicAdd(&global_buffer[0], local_pos);
|
||||
atomicAdd(&global_buffer[1], local_neg);
|
||||
atomicAdd(&global_buffer[2], local_num);
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void centernet_reg_loss_kernel(
|
||||
const float* __restrict__ pred,
|
||||
const float* __restrict__ gt,
|
||||
const float* __restrict__ mask,
|
||||
float* __restrict__ global_buffer,
|
||||
int output_idx,
|
||||
int n_elements,
|
||||
int spatial)
|
||||
{
|
||||
int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int stride = gridDim.x * blockDim.x;
|
||||
|
||||
float local_loss = 0.0f;
|
||||
int stride_pred = 2 * spatial;
|
||||
|
||||
for (int idx = tid; idx < n_elements; idx += stride) {
|
||||
int mask_idx = (idx / stride_pred) * spatial + (idx % spatial);
|
||||
|
||||
if (mask[mask_idx] == 1.0f) {
|
||||
local_loss += fabsf(pred[idx] - gt[idx]);
|
||||
}
|
||||
}
|
||||
|
||||
local_loss = warp_reduce(local_loss);
|
||||
|
||||
if ((threadIdx.x % 32) == 0) {
|
||||
atomicAdd(&global_buffer[output_idx], local_loss);
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void centernet_final_loss_kernel(
|
||||
float* __restrict__ buffer,
|
||||
float* __restrict__ out)
|
||||
{
|
||||
if (threadIdx.x == 0) {
|
||||
float pos_sum = buffer[0];
|
||||
float neg_sum = buffer[1];
|
||||
float num_pos = buffer[2];
|
||||
float wh_sum = buffer[3];
|
||||
float reg_sum = buffer[4];
|
||||
|
||||
float hm_loss = 0.0f;
|
||||
if (num_pos > 0.0f) {
|
||||
hm_loss = -(pos_sum + neg_sum) / num_pos;
|
||||
wh_sum /= num_pos;
|
||||
reg_sum /= num_pos;
|
||||
} else {
|
||||
hm_loss = -neg_sum;
|
||||
}
|
||||
|
||||
out[0] = hm_loss + 0.1f * wh_sum + 1.0f * reg_sum;
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor launch_centernet_loss(
|
||||
torch::Tensor pred_hm, torch::Tensor gt_hm,
|
||||
torch::Tensor pred_wh, torch::Tensor gt_wh,
|
||||
torch::Tensor pred_reg, torch::Tensor gt_reg,
|
||||
torch::Tensor mask)
|
||||
{
|
||||
auto options = pred_hm.options();
|
||||
auto out = torch::empty({1}, options);
|
||||
auto buffer = torch::zeros({5}, options);
|
||||
|
||||
int n_hm = pred_hm.numel();
|
||||
int threads = 256;
|
||||
int blocks_hm = (n_hm + threads - 1) / threads;
|
||||
if (blocks_hm > 256) blocks_hm = 256;
|
||||
|
||||
centernet_hm_loss_kernel<<<blocks_hm, threads>>>(
|
||||
pred_hm.data_ptr<float>(),
|
||||
gt_hm.data_ptr<float>(),
|
||||
buffer.data_ptr<float>(),
|
||||
n_hm
|
||||
);
|
||||
|
||||
int n_reg = pred_wh.numel();
|
||||
int batch_size = pred_wh.size(0);
|
||||
int spatial = pred_wh.size(2) * pred_wh.size(3);
|
||||
|
||||
int blocks_reg = (n_reg + threads - 1) / threads;
|
||||
if (blocks_reg > 256) blocks_reg = 256;
|
||||
|
||||
centernet_reg_loss_kernel<<<blocks_reg, threads>>>(
|
||||
pred_wh.data_ptr<float>(),
|
||||
gt_wh.data_ptr<float>(),
|
||||
mask.data_ptr<float>(),
|
||||
buffer.data_ptr<float>(),
|
||||
3,
|
||||
n_reg,
|
||||
spatial
|
||||
);
|
||||
|
||||
centernet_reg_loss_kernel<<<blocks_reg, threads>>>(
|
||||
pred_reg.data_ptr<float>(),
|
||||
gt_reg.data_ptr<float>(),
|
||||
mask.data_ptr<float>(),
|
||||
buffer.data_ptr<float>(),
|
||||
4,
|
||||
n_reg,
|
||||
spatial
|
||||
);
|
||||
|
||||
centernet_final_loss_kernel<<<1, 1>>>(
|
||||
buffer.data_ptr<float>(),
|
||||
out.data_ptr<float>()
|
||||
);
|
||||
|
||||
return out;
|
||||
}
|
||||
"""
|
||||
|
||||
cpp_source = """
|
||||
torch::Tensor launch_centernet_loss(
|
||||
torch::Tensor pred_hm, torch::Tensor gt_hm,
|
||||
torch::Tensor pred_wh, torch::Tensor gt_wh,
|
||||
torch::Tensor pred_reg, torch::Tensor gt_reg,
|
||||
torch::Tensor mask);
|
||||
"""
|
||||
|
||||
centernet_loss_module = load_inline(
|
||||
name='centernet_loss_opt',
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=cuda_source,
|
||||
functions=['launch_centernet_loss'],
|
||||
verbose=False
|
||||
)
|
||||
|
||||
|
||||
class ModelNew(nn.Module):
|
||||
def __init__(self):
|
||||
super(ModelNew, self).__init__()
|
||||
self.op = centernet_loss_module
|
||||
|
||||
def forward(self, pred_hm, gt_hm, pred_wh, gt_wh, pred_reg, gt_reg, mask):
|
||||
return self.op.launch_centernet_loss(
|
||||
pred_hm.contiguous(), gt_hm.contiguous(),
|
||||
pred_wh.contiguous(), gt_wh.contiguous(),
|
||||
pred_reg.contiguous(), gt_reg.contiguous(),
|
||||
mask.contiguous()
|
||||
)
|
||||
|
|
@ -0,0 +1,57 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self):
|
||||
super(Model, self).__init__()
|
||||
|
||||
def forward(self, pred_hm, gt_hm, pred_wh, gt_wh, pred_reg, gt_reg, mask):
|
||||
pred_hm = torch.clamp(pred_hm, 1e-6, 1 - 1e-6)
|
||||
|
||||
pos_inds = gt_hm.eq(1).float()
|
||||
neg_inds = gt_hm.lt(1).float()
|
||||
neg_weights = torch.pow(1 - gt_hm, 4)
|
||||
|
||||
pos_loss = torch.log(pred_hm) * torch.pow(1 - pred_hm, 2) * pos_inds
|
||||
neg_loss = torch.log(1 - pred_hm) * torch.pow(pred_hm, 2) * neg_weights * neg_inds
|
||||
|
||||
num_pos = pos_inds.sum()
|
||||
pos_loss_sum = pos_loss.sum()
|
||||
neg_loss_sum = neg_loss.sum()
|
||||
|
||||
if num_pos > 0:
|
||||
hm_loss = -(pos_loss_sum + neg_loss_sum) / num_pos
|
||||
else:
|
||||
hm_loss = -neg_loss_sum
|
||||
|
||||
mask_expanded = mask.expand_as(pred_wh)
|
||||
wh_loss = torch.sum(torch.abs(pred_wh - gt_wh) * mask_expanded)
|
||||
reg_loss = torch.sum(torch.abs(pred_reg - gt_reg) * mask_expanded)
|
||||
|
||||
if num_pos > 0:
|
||||
wh_loss = wh_loss / num_pos
|
||||
reg_loss = reg_loss / num_pos
|
||||
|
||||
return hm_loss + 0.1 * wh_loss + 1.0 * reg_loss
|
||||
|
||||
|
||||
batch_size = 4
|
||||
channels = 4
|
||||
height = 128
|
||||
width = 128
|
||||
|
||||
|
||||
def get_inputs():
|
||||
pred_hm = torch.sigmoid(torch.randn(batch_size, channels, height, width))
|
||||
gt_hm = torch.bernoulli(torch.full((batch_size, channels, height, width), 0.1))
|
||||
pred_wh = torch.randn(batch_size, 2, height, width)
|
||||
gt_wh = torch.randn(batch_size, 2, height, width)
|
||||
pred_reg = torch.randn(batch_size, 2, height, width)
|
||||
gt_reg = torch.randn(batch_size, 2, height, width)
|
||||
mask = torch.bernoulli(torch.full((batch_size, 1, height, width), 0.1))
|
||||
return [pred_hm, gt_hm, pred_wh, gt_wh, pred_reg, gt_reg, mask]
|
||||
|
||||
|
||||
def get_init_inputs():
|
||||
return []
|
||||
|
|
@ -0,0 +1,88 @@
|
|||
You write custom CUDA kernels to replace the pytorch operators in the given GeGLU 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 chunk+gelu+elementwise_mul), or algorithmic changes (such as optimized memory access patterns). You are only limited by your imagination.
|
||||
|
||||
This code implements CenterNet loss with CUDA optimizations:
|
||||
|
||||
Multi-kernel design - Separate kernels for heatmap loss and regression losses (width/height + offset).
|
||||
|
||||
Global buffer accumulation - Uses global memory buffer (5 floats) to accumulate partial sums from all threads.
|
||||
|
||||
Warp reduction + atomic addition - Combines warp-level reduction with atomic adds for thread-safe accumulation.
|
||||
|
||||
Custom focal loss variant - Implements modified focal loss with positive/negative term separation.
|
||||
|
||||
Masked regression loss - Only computes loss where mask=1 (valid positions).
|
||||
|
||||
Numerical stability - Clips predictions to [1e-6, 0.999999] to avoid log(0) issues.
|
||||
|
||||
Weighted negative samples - Applies (1-t)^4 weighting for negative samples in heatmap loss.
|
||||
|
||||
Grid-stride loops - Threads process multiple elements with stride for load balancing.
|
||||
|
||||
Final weighted combination - Combines heatmap loss + 0.1×WH loss + 1.0×reg loss in final kernel.
|
||||
|
||||
Memory efficiency - Reuses same kernel for both WH and regression losses with output index parameter.
|
||||
|
||||
Batch processing - Handles batched predictions with spatial dimensions.
|
||||
|
||||
|
||||
|
||||
Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is:
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self):
|
||||
super(Model, self).__init__()
|
||||
|
||||
def forward(self, pred_hm, gt_hm, pred_wh, gt_wh, pred_reg, gt_reg, mask):
|
||||
pred_hm = torch.clamp(pred_hm, 1e-6, 1 - 1e-6)
|
||||
|
||||
pos_inds = gt_hm.eq(1).float()
|
||||
neg_inds = gt_hm.lt(1).float()
|
||||
neg_weights = torch.pow(1 - gt_hm, 4)
|
||||
|
||||
pos_loss = torch.log(pred_hm) * torch.pow(1 - pred_hm, 2) * pos_inds
|
||||
neg_loss = torch.log(1 - pred_hm) * torch.pow(pred_hm, 2) * neg_weights * neg_inds
|
||||
|
||||
num_pos = pos_inds.sum()
|
||||
pos_loss_sum = pos_loss.sum()
|
||||
neg_loss_sum = neg_loss.sum()
|
||||
|
||||
if num_pos > 0:
|
||||
hm_loss = -(pos_loss_sum + neg_loss_sum) / num_pos
|
||||
else:
|
||||
hm_loss = -neg_loss_sum
|
||||
|
||||
mask_expanded = mask.expand_as(pred_wh)
|
||||
wh_loss = torch.sum(torch.abs(pred_wh - gt_wh) * mask_expanded)
|
||||
reg_loss = torch.sum(torch.abs(pred_reg - gt_reg) * mask_expanded)
|
||||
|
||||
if num_pos > 0:
|
||||
wh_loss = wh_loss / num_pos
|
||||
reg_loss = reg_loss / num_pos
|
||||
|
||||
return hm_loss + 0.1 * wh_loss + 1.0 * reg_loss
|
||||
|
||||
|
||||
batch_size = 4
|
||||
channels = 4
|
||||
height = 128
|
||||
width = 128
|
||||
|
||||
|
||||
def get_inputs():
|
||||
pred_hm = torch.sigmoid(torch.randn(batch_size, channels, height, width))
|
||||
gt_hm = torch.bernoulli(torch.full((batch_size, channels, height, width), 0.1))
|
||||
pred_wh = torch.randn(batch_size, 2, height, width)
|
||||
gt_wh = torch.randn(batch_size, 2, height, width)
|
||||
pred_reg = torch.randn(batch_size, 2, height, width)
|
||||
gt_reg = torch.randn(batch_size, 2, height, width)
|
||||
mask = torch.bernoulli(torch.full((batch_size, 1, height, width), 0.1))
|
||||
return [pred_hm, gt_hm, pred_wh, gt_wh, pred_reg, gt_reg, mask]
|
||||
|
||||
|
||||
def get_init_inputs():
|
||||
return []
|
||||
|
|
@ -0,0 +1,77 @@
|
|||
###########################################################
|
||||
# 性能和精度验证程序
|
||||
###########################################################
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import time
|
||||
from CenterNetLoss_torch import Model, get_inputs, get_init_inputs
|
||||
from CenterNetLoss_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