forked from ccf-ai-infra/GPUCodeForces
124 lines
4.4 KiB
Python
124 lines
4.4 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
from torch.utils.cpp_extension import load_inline
|
|
|
|
B, COORD = 32, 4
|
|
|
|
|
|
class ModelNew(nn.Module):
|
|
def __init__(self, block_size=256):
|
|
super().__init__()
|
|
self.block_size = block_size
|
|
self._compile_cuda_kernel()
|
|
|
|
def _compile_cuda_kernel(self):
|
|
cpp_source = """
|
|
#include <torch/extension.h>
|
|
torch::Tensor iou_fused_cuda(torch::Tensor pred, torch::Tensor target);
|
|
"""
|
|
|
|
cuda_source = f"""
|
|
#include <cuda_runtime.h>
|
|
#include <algorithm> // For std::min, std::max
|
|
#include <cmath> // For fmaxf, fminf
|
|
|
|
#define BLOCK_SIZE {self.block_size}
|
|
#define EPSILON 1e-6f
|
|
|
|
// 融合核函数：计算 IoU Loss (1 - IoU)
|
|
// 每个线程处理一个批次样本 (一个边界框对)
|
|
__global__ void iou_loss_fused_kernel(
|
|
const float* __restrict__ P, // B x 4 (x1, y1, x2, y2)
|
|
const float* __restrict__ T, // B x 4 (x1, y1, x2, y2)
|
|
float* __restrict__ D, // B x 1 (Loss per item)
|
|
int B,
|
|
int COORD
|
|
) {{
|
|
int b = blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (b < B) {{
|
|
|
|
// 预测框坐标 (P[b, 0]...P[b, 3])
|
|
const float px1 = P[b * COORD + 0];
|
|
const float py1 = P[b * COORD + 1];
|
|
const float px2 = P[b * COORD + 2];
|
|
const float py2 = P[b * COORD + 3];
|
|
|
|
// 目标框坐标 (T[b, 0]...T[b, 3])
|
|
const float tx1 = T[b * COORD + 0];
|
|
const float ty1 = T[b * COORD + 1];
|
|
const float tx2 = T[b * COORD + 2];
|
|
const float ty2 = T[b * COORD + 3];
|
|
|
|
// 1. 交集区域坐标
|
|
const float ix1 = fmaxf(px1, tx1);
|
|
const float iy1 = fmaxf(py1, ty1);
|
|
const float ix2 = fminf(px2, tx2);
|
|
const float iy2 = fminf(py2, ty2);
|
|
|
|
// 2. 交集区域边长和面积
|
|
const float iw = fmaxf(ix2 - ix1, 0.0f);
|
|
const float ih = fmaxf(iy2 - iy1, 0.0f);
|
|
const float intersection = iw * ih;
|
|
|
|
// 3. 预测框和目标框面积
|
|
const float area_p = (px2 - px1) * (py2 - py1);
|
|
const float area_t = (tx2 - tx1) * (ty2 - ty1);
|
|
|
|
// 4. 并集区域面积: A_union = A_p + A_t - A_intersection
|
|
const float union_area = area_p + area_t - intersection;
|
|
|
|
// 5. IoU 和 Loss
|
|
const float iou = intersection / (union_area + EPSILON);
|
|
const float loss = 1.0f - iou;
|
|
|
|
// 存储每个样本的损失
|
|
D[b] = loss;
|
|
}}
|
|
}}
|
|
|
|
torch::Tensor iou_fused_cuda(torch::Tensor pred, torch::Tensor target) {{
|
|
|
|
TORCH_CHECK(pred.is_cuda() && target.is_cuda(), "Inputs must be CUDA tensors");
|
|
TORCH_CHECK(pred.dim() == 2 && target.dim() == 2, "Inputs must be 2D (B, COORD)");
|
|
|
|
pred = pred.contiguous();
|
|
target = target.contiguous();
|
|
|
|
int B = pred.size(0);
|
|
int COORD = pred.size(1);
|
|
|
|
// 输出 D 的形状是 B x 1 (每批次样本的损失)
|
|
auto D = torch::empty({{B, 1}}, pred.options());
|
|
|
|
int blocks = std::min((B + BLOCK_SIZE - 1) / BLOCK_SIZE, 1024);
|
|
|
|
iou_loss_fused_kernel<<<blocks, BLOCK_SIZE>>>(
|
|
pred.data_ptr<float>(),
|
|
target.data_ptr<float>(),
|
|
D.data_ptr<float>(),
|
|
B,
|
|
COORD
|
|
);
|
|
|
|
return D;
|
|
}}
|
|
"""
|
|
|
|
self.op = load_inline(
|
|
name='iou_loss_opt',
|
|
cpp_sources=cpp_source,
|
|
cuda_sources=cuda_source,
|
|
functions=['iou_fused_cuda'],
|
|
extra_cuda_cflags=['-O3', '--use_fast_math'],
|
|
verbose=False
|
|
)
|
|
|
|
def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
|
if not pred.is_cuda: pred = pred.cuda()
|
|
if not target.is_cuda: target = target.cuda()
|
|
|
|
# 调用融合内核，返回每批次样本的损失 (B x 1)
|
|
per_item_loss = self.op.iou_fused_cuda(pred, target)
|
|
|
|
# 在 C++ 侧对 B x 1 结果求平均，实现最终的标量损失
|
|
return per_item_loss.mean() |