forked from ccf-ai-infra/GPUCodeForces
finish hillloss #52
This commit is contained in:
parent
f876a28ada
commit
23f48ac7ce
|
|
@ -0,0 +1,150 @@
|
|||
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 hill_loss_cuda_forward(
|
||||
const torch::Tensor& logits,
|
||||
const torch::Tensor& targets,
|
||||
float lambda_val,
|
||||
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;
|
||||
};
|
||||
|
||||
// Sigmoid
|
||||
__device__ __forceinline__ float fast_sigmoid(float x) {
|
||||
return 1.0f / (1.0f + __expf(-x));
|
||||
}
|
||||
|
||||
// Softplus (stable -log(sigmoid(x)))
|
||||
__device__ __forceinline__ float fast_softplus_neg(float x) {
|
||||
// return log(1 + exp(-x))
|
||||
// if -x is large, exp(-x) -> inf.
|
||||
// if x < -80, result is -x.
|
||||
// if x > 20, result is 0.
|
||||
// Standard Softplus: log(1 + exp(x))
|
||||
// We want log(1 + exp(-x)) = Softplus(-x)
|
||||
float neg_x = -x;
|
||||
if (neg_x > 20.0f) return neg_x;
|
||||
return logf(1.0f + __expf(neg_x));
|
||||
}
|
||||
|
||||
__device__ __forceinline__ float compute_hill_point(
|
||||
float logit, float target, float lambda_val)
|
||||
{
|
||||
if (target > 0.5f) {
|
||||
// Positive: BCE
|
||||
// -log(p) = -log(sigmoid(logit)) = softplus(-logit)
|
||||
return fast_softplus_neg(logit);
|
||||
} else {
|
||||
// Negative: Hill
|
||||
// (lambda - p) * p^2
|
||||
float p = fast_sigmoid(logit);
|
||||
return (lambda_val - p) * p * p;
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void hill_loss_kernel(
|
||||
float* __restrict__ output,
|
||||
const float* __restrict__ logits,
|
||||
const float* __restrict__ targets,
|
||||
const int n_elements,
|
||||
const float lambda_val)
|
||||
{
|
||||
const int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int vec_n = n_elements / 4;
|
||||
int stride = blockDim.x * gridDim.x;
|
||||
|
||||
for (int i = idx; i < vec_n; i += stride) {
|
||||
Float4 l_vec = reinterpret_cast<const Float4*>(logits)[i];
|
||||
Float4 t_vec = reinterpret_cast<const Float4*>(targets)[i];
|
||||
Float4 out_vec;
|
||||
|
||||
out_vec.x = compute_hill_point(l_vec.x, t_vec.x, lambda_val);
|
||||
out_vec.y = compute_hill_point(l_vec.y, t_vec.y, lambda_val);
|
||||
out_vec.z = compute_hill_point(l_vec.z, t_vec.z, lambda_val);
|
||||
out_vec.w = compute_hill_point(l_vec.w, t_vec.w, lambda_val);
|
||||
|
||||
reinterpret_cast<Float4*>(output)[i] = out_vec;
|
||||
}
|
||||
|
||||
int tail_start = vec_n * 4;
|
||||
int scalar_idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
if (scalar_idx < (n_elements - tail_start)) {
|
||||
int real_idx = tail_start + scalar_idx;
|
||||
output[real_idx] = compute_hill_point(
|
||||
logits[real_idx], targets[real_idx], lambda_val);
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor hill_loss_cuda_forward(
|
||||
const torch::Tensor& logits,
|
||||
const torch::Tensor& targets,
|
||||
float lambda_val,
|
||||
std::string reduction)
|
||||
{
|
||||
TORCH_CHECK(logits.is_cuda() && targets.is_cuda(), "Inputs must be CUDA");
|
||||
TORCH_CHECK(logits.is_contiguous() && targets.is_contiguous(), "Inputs must be contiguous");
|
||||
TORCH_CHECK(logits.numel() == targets.numel(), "Shapes must match");
|
||||
|
||||
const int n = logits.numel();
|
||||
auto output = torch::empty_like(logits);
|
||||
|
||||
const int vec_n = n / 4;
|
||||
const int grid_size = (vec_n + BLOCK_SIZE - 1) / BLOCK_SIZE;
|
||||
int final_grid = (grid_size < 1) ? 1 : grid_size;
|
||||
if (final_grid > 65535) final_grid = 65535;
|
||||
|
||||
hill_loss_kernel<<<final_grid, BLOCK_SIZE>>>(
|
||||
output.data_ptr<float>(),
|
||||
logits.data_ptr<float>(),
|
||||
targets.data_ptr<float>(),
|
||||
n,
|
||||
lambda_val
|
||||
);
|
||||
|
||||
if (reduction == "mean") {
|
||||
return output.mean();
|
||||
} else if (reduction == "sum") {
|
||||
return output.sum();
|
||||
}
|
||||
|
||||
return output;
|
||||
}
|
||||
"""
|
||||
|
||||
class ModelNew(nn.Module):
|
||||
def __init__(self, lambda_val=1.5, reduction='none'):
|
||||
super(ModelNew, self).__init__()
|
||||
self.lambda_val = lambda_val
|
||||
self.reduction = reduction
|
||||
self.op = load_inline(
|
||||
name='hill_loss_op',
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=cuda_source,
|
||||
functions=['hill_loss_cuda_forward'],
|
||||
verbose=False,
|
||||
extra_cuda_cflags=['-O3']
|
||||
)
|
||||
|
||||
def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
|
||||
return self.op.hill_loss_cuda_forward(
|
||||
logits.contiguous(),
|
||||
targets.contiguous(),
|
||||
self.lambda_val,
|
||||
self.reduction
|
||||
)
|
||||
|
|
@ -0,0 +1,62 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
BATCH_SIZE = 128
|
||||
NUM_CLASSES = 80
|
||||
SHAPE = (BATCH_SIZE, NUM_CLASSES, 64, 64)
|
||||
|
||||
LAMBDA = 1.5
|
||||
REDUCTION = 'none'
|
||||
|
||||
class HillLoss(nn.Module):
|
||||
"""
|
||||
Hill Loss for Multi-Label Learning with Missing Labels.
|
||||
https://arxiv.org/pdf/2112.07368
|
||||
Negatives: (lambda - p) * p^2
|
||||
Positives: BCE
|
||||
"""
|
||||
def __init__(self, lambda_val=1.5, reduction='mean'):
|
||||
super(HillLoss, self).__init__()
|
||||
self.lambda_val = lambda_val
|
||||
self.reduction = reduction
|
||||
|
||||
def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
|
||||
# logits: arbitrary shape
|
||||
# targets: same shape, 0 or 1
|
||||
|
||||
probs = torch.sigmoid(logits)
|
||||
|
||||
# Positive Loss
|
||||
pos_loss = F.softplus(-logits)
|
||||
|
||||
# Negative Loss (Hill Loss)
|
||||
neg_loss = (self.lambda_val - probs) * (probs ** 2)
|
||||
|
||||
# Combine
|
||||
# loss = y * pos + (1-y) * neg
|
||||
loss = targets * pos_loss + (1 - targets) * neg_loss
|
||||
|
||||
if self.reduction == 'mean':
|
||||
return loss.mean()
|
||||
elif self.reduction == 'sum':
|
||||
return loss.sum()
|
||||
return loss
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, lambda_val=1.5, reduction='none'):
|
||||
super(Model, self).__init__()
|
||||
self.loss_fn = HillLoss(lambda_val=lambda_val, reduction=reduction)
|
||||
|
||||
def forward(self, logits, targets):
|
||||
return self.loss_fn(logits, targets)
|
||||
|
||||
def get_inputs():
|
||||
logits = torch.randn(SHAPE, dtype=torch.float32)
|
||||
# 稀疏标签，大部分为 0
|
||||
targets = (torch.rand(SHAPE) > 0.9).float()
|
||||
return [logits.contiguous(), targets.contiguous()]
|
||||
|
||||
def get_init_inputs():
|
||||
return [LAMBDA, REDUCTION]
|
||||
|
|
@ -0,0 +1,93 @@
|
|||
Write a custom CUDA kernel to optimize `Hill Loss` (from "Simple and Robust Loss Design for Multi-Label Learning with Missing Labels").
|
||||
|
||||
Formula:
|
||||
Loss = y * L_pos + (1 - y) * L_neg
|
||||
L_pos = -log(p) (Standard BCE for positives)
|
||||
L_neg = (lambda - p) * p^2 (Hill Loss for negatives)
|
||||
Where p = sigmoid(logit).
|
||||
Lambda is a hyperparameter (default 1.5).
|
||||
|
||||
Problem Analysis:
|
||||
1. Memory Efficiency: Standard implementation requires calculating sigmoid, creating masks for positive/negative samples, computing different loss branches, and merging. This generates intermediate tensors and multiple memory passes.
|
||||
2. Element-wise Fusion: The operation is strictly element-wise and computationally lightweight.
|
||||
|
||||
Optimization Strategy: Fused Element-wise Kernel with Vectorization
|
||||
|
||||
1. Flattened Input: Treat input tensors of any shape `(N, C, ...)` as a 1D array.
|
||||
|
||||
2. Vectorized Loads (float4): Use `float4` to load 4 logits and 4 targets at a time.
|
||||
|
||||
3. In-Register Logic:
|
||||
- Compute `p = sigmoid(logit)`.
|
||||
- Branchless or If-Else logic based on `target`:
|
||||
- If `target == 1`: `loss = -log(p)` (using stable `log_sigmoid` equivalent).
|
||||
- If `target == 0`: `loss = (lambda - p) * p * p`.
|
||||
- Store result directly.
|
||||
|
||||
4. Numerical Stability: Ensure `log(p)` handles edge cases safely, or use `log_sigmoid(logit)` for positives.
|
||||
|
||||
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 = 128
|
||||
NUM_CLASSES = 80
|
||||
SHAPE = (BATCH_SIZE, NUM_CLASSES, 64, 64)
|
||||
|
||||
LAMBDA = 1.5
|
||||
REDUCTION = 'none'
|
||||
|
||||
class HillLoss(nn.Module):
|
||||
"""
|
||||
Hill Loss for Multi-Label Learning with Missing Labels.
|
||||
https://arxiv.org/pdf/2112.07368
|
||||
Negatives: (lambda - p) * p^2
|
||||
Positives: BCE
|
||||
"""
|
||||
def __init__(self, lambda_val=1.5, reduction='mean'):
|
||||
super(HillLoss, self).__init__()
|
||||
self.lambda_val = lambda_val
|
||||
self.reduction = reduction
|
||||
|
||||
def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
|
||||
# logits: arbitrary shape
|
||||
# targets: same shape, 0 or 1
|
||||
|
||||
probs = torch.sigmoid(logits)
|
||||
|
||||
# Positive Loss
|
||||
pos_loss = F.softplus(-logits)
|
||||
|
||||
# Negative Loss (Hill Loss)
|
||||
neg_loss = (self.lambda_val - probs) * (probs ** 2)
|
||||
|
||||
# Combine
|
||||
# loss = y * pos + (1-y) * neg
|
||||
loss = targets * pos_loss + (1 - targets) * neg_loss
|
||||
|
||||
if self.reduction == 'mean':
|
||||
return loss.mean()
|
||||
elif self.reduction == 'sum':
|
||||
return loss.sum()
|
||||
return loss
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, lambda_val=1.5, reduction='none'):
|
||||
super(Model, self).__init__()
|
||||
self.loss_fn = HillLoss(lambda_val=lambda_val, reduction=reduction)
|
||||
|
||||
def forward(self, logits, targets):
|
||||
return self.loss_fn(logits, targets)
|
||||
|
||||
def get_inputs():
|
||||
logits = torch.randn(SHAPE, dtype=torch.float32)
|
||||
# 稀疏标签，大部分为 0
|
||||
targets = (torch.rand(SHAPE) > 0.9).float()
|
||||
return [logits.contiguous(), targets.contiguous()]
|
||||
|
||||
def get_init_inputs():
|
||||
return [LAMBDA, REDUCTION]
|
||||
|
|
@ -0,0 +1,74 @@
|
|||
###########################################################
|
||||
# 性能和精度验证程序
|
||||
###########################################################
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import time
|
||||
from hillloss_torch import Model,get_inputs,get_init_inputs
|
||||
from hillloss_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