forked from ccf-ai-infra/GPUCodeForces
finish DIoUloss #54
This commit is contained in:
parent
f876a28ada
commit
8e18d22b2e
|
|
@ -0,0 +1,132 @@
|
|||
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 diou_loss_cuda_forward(
|
||||
const torch::Tensor& b1,
|
||||
const torch::Tensor& b2,
|
||||
float eps,
|
||||
std::string reduction);
|
||||
"""
|
||||
|
||||
cuda_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <math.h>
|
||||
|
||||
#define BLOCK_SIZE 256
|
||||
|
||||
struct __align__(16) Float4 {
|
||||
float x, y, z, w;
|
||||
};
|
||||
|
||||
__global__ void diou_loss_kernel(
|
||||
float* __restrict__ output,
|
||||
const float* __restrict__ b1,
|
||||
const float* __restrict__ b2,
|
||||
const int num_boxes,
|
||||
const float eps)
|
||||
{
|
||||
const int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx >= num_boxes) return;
|
||||
|
||||
// Vectorized Load
|
||||
Float4 box1 = reinterpret_cast<const Float4*>(b1)[idx];
|
||||
Float4 box2 = reinterpret_cast<const Float4*>(b2)[idx];
|
||||
|
||||
float b1_x1 = box1.x; float b1_y1 = box1.y; float b1_x2 = box1.z; float b1_y2 = box1.w;
|
||||
float b2_x1 = box2.x; float b2_y1 = box2.y; float b2_x2 = box2.z; float b2_y2 = box2.w;
|
||||
|
||||
// Compute Intersection
|
||||
float inter_x1 = fmaxf(b1_x1, b2_x1);
|
||||
float inter_y1 = fmaxf(b1_y1, b2_y1);
|
||||
float inter_x2 = fminf(b1_x2, b2_x2);
|
||||
float inter_y2 = fminf(b1_y2, b2_y2);
|
||||
|
||||
float inter_w = fmaxf(inter_x2 - inter_x1, 0.0f);
|
||||
float inter_h = fmaxf(inter_y2 - inter_y1, 0.0f);
|
||||
float inter_area = inter_w * inter_h;
|
||||
|
||||
// Compute Union
|
||||
float area1 = (b1_x2 - b1_x1) * (b1_y2 - b1_y1);
|
||||
float area2 = (b2_x2 - b2_x1) * (b2_y2 - b2_y1);
|
||||
float union_area = area1 + area2 - inter_area + eps;
|
||||
|
||||
// IoU
|
||||
float iou = inter_area / union_area;
|
||||
|
||||
// Compute Centers
|
||||
float c1_x = (b1_x1 + b1_x2) * 0.5f;
|
||||
float c1_y = (b1_y1 + b1_y2) * 0.5f;
|
||||
float c2_x = (b2_x1 + b2_x2) * 0.5f;
|
||||
float c2_y = (b2_y1 + b2_y2) * 0.5f;
|
||||
|
||||
float center_dist_sq = (c1_x - c2_x)*(c1_x - c2_x) + (c1_y - c2_y)*(c1_y - c2_y);
|
||||
|
||||
// Compute Enclosing Diagonal
|
||||
float enc_x1 = fminf(b1_x1, b2_x1);
|
||||
float enc_y1 = fminf(b1_y1, b2_y1);
|
||||
float enc_x2 = fmaxf(b1_x2, b2_x2);
|
||||
float enc_y2 = fmaxf(b1_y2, b2_y2);
|
||||
|
||||
float diag_dist_sq = (enc_x2 - enc_x1)*(enc_x2 - enc_x1) + (enc_y2 - enc_y1)*(enc_y2 - enc_y1) + eps;
|
||||
|
||||
// DIoU
|
||||
float diou = iou - center_dist_sq / diag_dist_sq;
|
||||
|
||||
output[idx] = 1.0f - diou;
|
||||
}
|
||||
|
||||
torch::Tensor diou_loss_cuda_forward(
|
||||
const torch::Tensor& b1,
|
||||
const torch::Tensor& b2,
|
||||
float eps,
|
||||
std::string reduction)
|
||||
{
|
||||
TORCH_CHECK(b1.is_cuda() && b2.is_cuda(), "Inputs must be CUDA");
|
||||
TORCH_CHECK(b1.is_contiguous() && b2.is_contiguous(), "Inputs must be contiguous");
|
||||
TORCH_CHECK(b1.size(1) == 4 && b2.size(1) == 4, "Boxes must be (N, 4)");
|
||||
|
||||
int n = b1.size(0);
|
||||
auto output = torch::empty({n}, b1.options());
|
||||
|
||||
int grid_size = (n + BLOCK_SIZE - 1) / BLOCK_SIZE;
|
||||
|
||||
diou_loss_kernel<<<grid_size, BLOCK_SIZE>>>(
|
||||
output.data_ptr<float>(),
|
||||
b1.data_ptr<float>(),
|
||||
b2.data_ptr<float>(),
|
||||
n,
|
||||
eps
|
||||
);
|
||||
|
||||
if (reduction == "mean") {
|
||||
return output.mean();
|
||||
} else if (reduction == "sum") {
|
||||
return output.sum();
|
||||
}
|
||||
|
||||
return output;
|
||||
}
|
||||
"""
|
||||
|
||||
class ModelNew(nn.Module):
|
||||
def __init__(self, eps=1e-7, reduction='none'):
|
||||
super(ModelNew, self).__init__()
|
||||
self.eps = eps
|
||||
self.reduction = reduction
|
||||
self.op = load_inline(
|
||||
name='diou_loss_op',
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=cuda_source,
|
||||
functions=['diou_loss_cuda_forward'],
|
||||
verbose=False,
|
||||
extra_cuda_cflags=['-O3']
|
||||
)
|
||||
|
||||
def forward(self, b1: torch.Tensor, b2: torch.Tensor) -> torch.Tensor:
|
||||
return self.op.diou_loss_cuda_forward(b1.contiguous(), b2.contiguous(), self.eps, self.reduction)
|
||||
|
|
@ -0,0 +1,99 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
import math
|
||||
|
||||
BATCH_SIZE = 128
|
||||
NUM_ANCHORS = 8400
|
||||
TOTAL_BOXES = BATCH_SIZE * NUM_ANCHORS
|
||||
SHAPE = (TOTAL_BOXES, 4)
|
||||
|
||||
EPS = 1e-7
|
||||
REDUCTION = 'none'
|
||||
|
||||
class DIoULoss(nn.Module):
|
||||
"""
|
||||
Distance-IoU Loss (AAAI 2020)
|
||||
https://arxiv.org/pdf/1911.08287
|
||||
"""
|
||||
def __init__(self, eps=1e-7, reduction='mean'):
|
||||
super(DIoULoss, self).__init__()
|
||||
self.eps = eps
|
||||
self.reduction = reduction
|
||||
|
||||
def forward(self, b1: torch.Tensor, b2: torch.Tensor) -> torch.Tensor:
|
||||
# b1, b2: (N, 4) -> (x1, y1, x2, y2)
|
||||
|
||||
# Coordinate Parsing
|
||||
b1_x1, b1_y1, b1_x2, b1_y2 = b1[:, 0], b1[:, 1], b1[:, 2], b1[:, 3]
|
||||
b2_x1, b2_y1, b2_x2, b2_y2 = b2[:, 0], b2[:, 1], b2[:, 2], b2[:, 3]
|
||||
|
||||
# IoU Calculation
|
||||
# Intersection
|
||||
inter_x1 = torch.max(b1_x1, b2_x1)
|
||||
inter_y1 = torch.max(b1_y1, b2_y1)
|
||||
inter_x2 = torch.min(b1_x2, b2_x2)
|
||||
inter_y2 = torch.min(b1_y2, b2_y2)
|
||||
|
||||
inter_w = (inter_x2 - inter_x1).clamp(min=0)
|
||||
inter_h = (inter_y2 - inter_y1).clamp(min=0)
|
||||
inter_area = inter_w * inter_h
|
||||
|
||||
# Union
|
||||
w1, h1 = b1_x2 - b1_x1, b1_y2 - b1_y1
|
||||
w2, h2 = b2_x2 - b2_x1, b2_y2 - b2_y1
|
||||
union_area = w1 * h1 + w2 * h2 - inter_area + self.eps
|
||||
|
||||
iou = inter_area / union_area
|
||||
|
||||
# DIoU Term
|
||||
# Center points
|
||||
c1_x, c1_y = (b1_x1 + b1_x2) / 2, (b1_y1 + b1_y2) / 2
|
||||
c2_x, c2_y = (b2_x1 + b2_x2) / 2, (b2_y1 + b2_y2) / 2
|
||||
|
||||
# Center distance squared (rho^2)
|
||||
center_dist_sq = (c1_x - c2_x)**2 + (c1_y - c2_y)**2
|
||||
|
||||
# Enclosing box
|
||||
enc_x1 = torch.min(b1_x1, b2_x1)
|
||||
enc_y1 = torch.min(b1_y1, b2_y1)
|
||||
enc_x2 = torch.max(b1_x2, b2_x2)
|
||||
enc_y2 = torch.max(b1_y2, b2_y2)
|
||||
|
||||
# Diagonal squared (c^2)
|
||||
diag_dist_sq = (enc_x2 - enc_x1)**2 + (enc_y2 - enc_y1)**2 + self.eps
|
||||
|
||||
# Final Loss
|
||||
diou = iou - center_dist_sq / diag_dist_sq
|
||||
loss = 1.0 - diou
|
||||
|
||||
if self.reduction == 'mean':
|
||||
return loss.mean()
|
||||
elif self.reduction == 'sum':
|
||||
return loss.sum()
|
||||
return loss
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, eps=1e-7, reduction='none'):
|
||||
super(Model, self).__init__()
|
||||
self.loss_fn = DIoULoss(eps=eps, reduction=reduction)
|
||||
|
||||
def forward(self, b1, b2):
|
||||
return self.loss_fn(b1, b2)
|
||||
|
||||
def get_inputs():
|
||||
b1 = torch.randn(SHAPE, dtype=torch.float32)
|
||||
b2 = torch.randn(SHAPE, dtype=torch.float32)
|
||||
|
||||
# Make valid (x2 > x1)
|
||||
def make_valid(b):
|
||||
x_min, _ = torch.min(b[:, [0, 2]], dim=1, keepdim=True)
|
||||
x_max, _ = torch.max(b[:, [0, 2]], dim=1, keepdim=True)
|
||||
y_min, _ = torch.min(b[:, [1, 3]], dim=1, keepdim=True)
|
||||
y_max, _ = torch.max(b[:, [1, 3]], dim=1, keepdim=True)
|
||||
# Ensure some overlap potential
|
||||
return torch.cat([x_min, y_min, x_max + 1.0, y_max + 1.0], dim=1).contiguous()
|
||||
|
||||
return [make_valid(b1), make_valid(b2)]
|
||||
|
||||
def get_init_inputs():
|
||||
return [EPS, REDUCTION]
|
||||
|
|
@ -0,0 +1,128 @@
|
|||
Write a custom CUDA kernel to optimize `DIoU Loss` (Distance-IoU Loss).
|
||||
|
||||
Formula: Loss = 1 - IoU + (distance_centers^2 / diagonal_enclosing^2)
|
||||
Where:
|
||||
- IoU is Intersection over Union.
|
||||
- distance_centers is the Euclidean distance between the center points of the two boxes.
|
||||
- diagonal_enclosing is the diagonal length of the smallest enclosing box covering both boxes.
|
||||
|
||||
Problem Analysis:
|
||||
1. Geometric Computations: Requires calculating centers, intersection area, union area, and enclosing box dimensions for every pair. Standard implementation creates multiple intermediate tensors.
|
||||
2. Memory Efficiency: Fusing these operations into a single kernel reduces global memory traffic significantly.
|
||||
|
||||
Optimization Strategy: Fused Element-wise Kernel with Vectorization
|
||||
|
||||
1. Input Format: Boxes are typically (x1, y1, x2, y2). Treat input as (N, 4).
|
||||
2. Vectorized Loads (float4): Load an entire box (4 floats) into registers using a single 128-bit instruction.
|
||||
3. In-Register Logic:
|
||||
- Compute Area1, Area2.
|
||||
- Compute Intersection (x1_max, y1_max, x2_min, y2_min).
|
||||
- Compute Union.
|
||||
- Compute Center points (ctx, cty) for both boxes.
|
||||
- Compute Enclosing Box (x1_min, y1_min, x2_max, y2_max) and its diagonal squared.
|
||||
- Compute center distance squared.
|
||||
- Combine to get DIoU Loss.
|
||||
4. Numerical Stability: Add epsilon to denominators (Union area and Diagonal squared).
|
||||
|
||||
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 math
|
||||
|
||||
BATCH_SIZE = 128
|
||||
NUM_ANCHORS = 8400
|
||||
TOTAL_BOXES = BATCH_SIZE * NUM_ANCHORS
|
||||
SHAPE = (TOTAL_BOXES, 4)
|
||||
|
||||
EPS = 1e-7
|
||||
REDUCTION = 'none'
|
||||
|
||||
class DIoULoss(nn.Module):
|
||||
"""
|
||||
Distance-IoU Loss (AAAI 2020)
|
||||
https://arxiv.org/pdf/1911.08287
|
||||
"""
|
||||
def __init__(self, eps=1e-7, reduction='mean'):
|
||||
super(DIoULoss, self).__init__()
|
||||
self.eps = eps
|
||||
self.reduction = reduction
|
||||
|
||||
def forward(self, b1: torch.Tensor, b2: torch.Tensor) -> torch.Tensor:
|
||||
# b1, b2: (N, 4) -> (x1, y1, x2, y2)
|
||||
|
||||
# Coordinate Parsing
|
||||
b1_x1, b1_y1, b1_x2, b1_y2 = b1[:, 0], b1[:, 1], b1[:, 2], b1[:, 3]
|
||||
b2_x1, b2_y1, b2_x2, b2_y2 = b2[:, 0], b2[:, 1], b2[:, 2], b2[:, 3]
|
||||
|
||||
# IoU Calculation
|
||||
# Intersection
|
||||
inter_x1 = torch.max(b1_x1, b2_x1)
|
||||
inter_y1 = torch.max(b1_y1, b2_y1)
|
||||
inter_x2 = torch.min(b1_x2, b2_x2)
|
||||
inter_y2 = torch.min(b1_y2, b2_y2)
|
||||
|
||||
inter_w = (inter_x2 - inter_x1).clamp(min=0)
|
||||
inter_h = (inter_y2 - inter_y1).clamp(min=0)
|
||||
inter_area = inter_w * inter_h
|
||||
|
||||
# Union
|
||||
w1, h1 = b1_x2 - b1_x1, b1_y2 - b1_y1
|
||||
w2, h2 = b2_x2 - b2_x1, b2_y2 - b2_y1
|
||||
union_area = w1 * h1 + w2 * h2 - inter_area + self.eps
|
||||
|
||||
iou = inter_area / union_area
|
||||
|
||||
# DIoU Term
|
||||
# Center points
|
||||
c1_x, c1_y = (b1_x1 + b1_x2) / 2, (b1_y1 + b1_y2) / 2
|
||||
c2_x, c2_y = (b2_x1 + b2_x2) / 2, (b2_y1 + b2_y2) / 2
|
||||
|
||||
# Center distance squared (rho^2)
|
||||
center_dist_sq = (c1_x - c2_x)**2 + (c1_y - c2_y)**2
|
||||
|
||||
# Enclosing box
|
||||
enc_x1 = torch.min(b1_x1, b2_x1)
|
||||
enc_y1 = torch.min(b1_y1, b2_y1)
|
||||
enc_x2 = torch.max(b1_x2, b2_x2)
|
||||
enc_y2 = torch.max(b1_y2, b2_y2)
|
||||
|
||||
# Diagonal squared (c^2)
|
||||
diag_dist_sq = (enc_x2 - enc_x1)**2 + (enc_y2 - enc_y1)**2 + self.eps
|
||||
|
||||
# Final Loss
|
||||
diou = iou - center_dist_sq / diag_dist_sq
|
||||
loss = 1.0 - diou
|
||||
|
||||
if self.reduction == 'mean':
|
||||
return loss.mean()
|
||||
elif self.reduction == 'sum':
|
||||
return loss.sum()
|
||||
return loss
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, eps=1e-7, reduction='none'):
|
||||
super(Model, self).__init__()
|
||||
self.loss_fn = DIoULoss(eps=eps, reduction=reduction)
|
||||
|
||||
def forward(self, b1, b2):
|
||||
return self.loss_fn(b1, b2)
|
||||
|
||||
def get_inputs():
|
||||
b1 = torch.randn(SHAPE, dtype=torch.float32)
|
||||
b2 = torch.randn(SHAPE, dtype=torch.float32)
|
||||
|
||||
# Make valid (x2 > x1)
|
||||
def make_valid(b):
|
||||
x_min, _ = torch.min(b[:, [0, 2]], dim=1, keepdim=True)
|
||||
x_max, _ = torch.max(b[:, [0, 2]], dim=1, keepdim=True)
|
||||
y_min, _ = torch.min(b[:, [1, 3]], dim=1, keepdim=True)
|
||||
y_max, _ = torch.max(b[:, [1, 3]], dim=1, keepdim=True)
|
||||
# Ensure some overlap potential
|
||||
return torch.cat([x_min, y_min, x_max + 1.0, y_max + 1.0], dim=1).contiguous()
|
||||
|
||||
return [make_valid(b1), make_valid(b2)]
|
||||
|
||||
def get_init_inputs():
|
||||
return [EPS, REDUCTION]
|
||||
|
|
@ -0,0 +1,74 @@
|
|||
###########################################################
|
||||
# 性能和精度验证程序
|
||||
###########################################################
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import time
|
||||
from DIoUloss_torch import Model,get_inputs,get_init_inputs
|
||||
from DIoUloss_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