forked from ccf-ai-infra/GPUCodeForces
parent
f989885dde
commit
992a67c865
|
|
@ -0,0 +1,159 @@
|
|||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
canberra_adaptive_triplet_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <math.h>
|
||||
|
||||
#define BLOCK_SIZE 256
|
||||
|
||||
__device__ float sigmoidf(float x) {
|
||||
return 1.0f / (1.0f + expf(-x));
|
||||
}
|
||||
|
||||
__global__ void canberra_adaptive_triplet_loss_kernel(
|
||||
const float* __restrict__ anchor,
|
||||
const float* __restrict__ positive,
|
||||
const float* __restrict__ negative,
|
||||
float* __restrict__ losses,
|
||||
int feature_dim,
|
||||
float base_margin,
|
||||
float alpha,
|
||||
float beta
|
||||
) {
|
||||
extern __shared__ float sdata[];
|
||||
|
||||
int tid = threadIdx.x;
|
||||
int sample_idx = blockIdx.x;
|
||||
|
||||
float pos_partial_sum = 0.0f;
|
||||
float neg_partial_sum = 0.0f;
|
||||
int base_addr = sample_idx * feature_dim;
|
||||
|
||||
// Each thread processes multiple elements with stride
|
||||
for (int i = tid; i < feature_dim; i += blockDim.x) {
|
||||
// Load all three values at once for better memory coalescing
|
||||
float anchor_val = anchor[base_addr + i];
|
||||
float pos_val = positive[base_addr + i];
|
||||
float neg_val = negative[base_addr + i];
|
||||
|
||||
// Compute anchor-positive Canberra distance
|
||||
float pos_abs_diff = fabsf(anchor_val - pos_val);
|
||||
float pos_abs_sum = fabsf(anchor_val) + fabsf(pos_val);
|
||||
if (pos_abs_sum > 0.0f) {
|
||||
pos_partial_sum += pos_abs_diff / pos_abs_sum;
|
||||
}
|
||||
|
||||
// Compute anchor-negative Canberra distance
|
||||
float neg_abs_diff = fabsf(anchor_val - neg_val);
|
||||
float neg_abs_sum = fabsf(anchor_val) + fabsf(neg_val);
|
||||
if (neg_abs_sum > 0.0f) {
|
||||
neg_partial_sum += neg_abs_diff / neg_abs_sum;
|
||||
}
|
||||
}
|
||||
|
||||
// Store partial sums in shared memory
|
||||
sdata[tid] = pos_partial_sum;
|
||||
sdata[tid + blockDim.x] = neg_partial_sum;
|
||||
__syncthreads();
|
||||
|
||||
// Parallel reduction for both distances
|
||||
for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
|
||||
if (tid < stride) {
|
||||
sdata[tid] += sdata[tid + stride];
|
||||
sdata[tid + blockDim.x] += sdata[tid + blockDim.x + stride];
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
// First thread computes adaptive triplet loss
|
||||
if (tid == 0) {
|
||||
float pos_dist = sdata[0];
|
||||
float neg_dist = sdata[blockDim.x];
|
||||
|
||||
// Compute adaptive margin
|
||||
float difficulty = sigmoidf(beta * (pos_dist - neg_dist));
|
||||
float adaptive_margin = base_margin + alpha * difficulty;
|
||||
|
||||
// Adaptive triplet loss
|
||||
float triplet_loss = pos_dist - neg_dist + adaptive_margin;
|
||||
losses[sample_idx] = fmaxf(triplet_loss, 0.0f);
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor canberra_adaptive_triplet_loss_cuda(
|
||||
torch::Tensor anchor,
|
||||
torch::Tensor positive,
|
||||
torch::Tensor negative,
|
||||
float base_margin,
|
||||
float alpha,
|
||||
float beta
|
||||
) {
|
||||
TORCH_CHECK(anchor.scalar_type() == torch::kFloat32, "Anchor must be float32");
|
||||
TORCH_CHECK(positive.scalar_type() == torch::kFloat32, "Positive must be float32");
|
||||
TORCH_CHECK(negative.scalar_type() == torch::kFloat32, "Negative must be float32");
|
||||
TORCH_CHECK(anchor.sizes() == positive.sizes(), "Anchor and positive must have the same shape");
|
||||
TORCH_CHECK(anchor.sizes() == negative.sizes(), "Anchor and negative must have the same shape");
|
||||
|
||||
auto anchor_contig = anchor.contiguous();
|
||||
auto positive_contig = positive.contiguous();
|
||||
auto negative_contig = negative.contiguous();
|
||||
|
||||
int batch_size = anchor_contig.size(0);
|
||||
int feature_dim = anchor_contig.size(1);
|
||||
|
||||
auto losses = torch::zeros({batch_size}, anchor.options());
|
||||
|
||||
const int block_size = BLOCK_SIZE;
|
||||
size_t shared_mem = 2 * block_size * sizeof(float); // For both positive and negative partial sums
|
||||
|
||||
canberra_adaptive_triplet_loss_kernel<<<batch_size, block_size, shared_mem>>>(
|
||||
anchor_contig.data_ptr<float>(),
|
||||
positive_contig.data_ptr<float>(),
|
||||
negative_contig.data_ptr<float>(),
|
||||
losses.data_ptr<float>(),
|
||||
feature_dim,
|
||||
base_margin,
|
||||
alpha,
|
||||
beta
|
||||
);
|
||||
|
||||
return losses;
|
||||
}
|
||||
"""
|
||||
|
||||
canberra_adaptive_triplet_cpp_source = """
|
||||
torch::Tensor canberra_adaptive_triplet_loss_cuda(
|
||||
torch::Tensor anchor,
|
||||
torch::Tensor positive,
|
||||
torch::Tensor negative,
|
||||
float base_margin,
|
||||
float alpha,
|
||||
float beta
|
||||
);
|
||||
"""
|
||||
|
||||
canberra_adaptive_triplet = load_inline(
|
||||
name="canberra_adaptive_triplet_loss",
|
||||
cpp_sources=canberra_adaptive_triplet_cpp_source,
|
||||
cuda_sources=canberra_adaptive_triplet_source,
|
||||
functions=["canberra_adaptive_triplet_loss_cuda"],
|
||||
extra_cuda_cflags=["-O3"], # 避免使用--use_fast_math
|
||||
verbose=True
|
||||
)
|
||||
|
||||
class ModelNew(torch.nn.Module):
|
||||
def __init__(self, base_margin=1.0, alpha=0.5, beta=0.1):
|
||||
super(ModelNew, self).__init__()
|
||||
self.base_margin = base_margin
|
||||
self.alpha = alpha
|
||||
self.beta = beta
|
||||
self.canberra_adaptive_triplet = canberra_adaptive_triplet
|
||||
|
||||
def forward(self, anchor, positive, negative):
|
||||
result = self.canberra_adaptive_triplet.canberra_adaptive_triplet_loss_cuda(
|
||||
anchor, positive, negative, self.base_margin, self.alpha, self.beta
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
return result
|
||||
|
|
@ -0,0 +1,68 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
Canberra Adaptive Triplet Loss implementation.
|
||||
Computes adaptive triplet loss using Canberra distance with dynamic margin adjustment.
|
||||
"""
|
||||
def __init__(self, base_margin=1.0, alpha=0.5, beta=0.1):
|
||||
super(Model, self).__init__()
|
||||
self.base_margin = base_margin
|
||||
self.alpha = alpha # Controls margin scaling
|
||||
self.beta = beta # Controls hardness sensitivity
|
||||
|
||||
def forward(self, anchor: torch.Tensor, positive: torch.Tensor, negative: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Compute Canberra Adaptive Triplet Loss.
|
||||
|
||||
Args:
|
||||
anchor (torch.Tensor): Anchor samples [batch_size, feature_dim]
|
||||
positive (torch.Tensor): Positive samples [batch_size, feature_dim]
|
||||
negative (torch.Tensor): Negative samples [batch_size, feature_dim]
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Adaptive triplet loss [batch_size]
|
||||
"""
|
||||
# Input validation
|
||||
if anchor.shape != positive.shape or anchor.shape != negative.shape:
|
||||
raise ValueError(f"All input tensors must have the same shape")
|
||||
|
||||
if anchor.dim() != 2:
|
||||
raise ValueError(f"Input tensors must be 2D, got {anchor.dim()}D")
|
||||
|
||||
# Compute Canberra distances using highly optimized built-in functions
|
||||
# Fuse operations to minimize memory allocation
|
||||
pos_dist = torch.sum(torch.where(
|
||||
(torch.abs(anchor) + torch.abs(positive)) > 0,
|
||||
torch.abs(anchor - positive) / (torch.abs(anchor) + torch.abs(positive)),
|
||||
torch.zeros_like(anchor)
|
||||
), dim=1)
|
||||
|
||||
neg_dist = torch.sum(torch.where(
|
||||
(torch.abs(anchor) + torch.abs(negative)) > 0,
|
||||
torch.abs(anchor - negative) / (torch.abs(anchor) + torch.abs(negative)),
|
||||
torch.zeros_like(anchor)
|
||||
), dim=1)
|
||||
|
||||
# Compute adaptive margin using built-in functions
|
||||
difficulty = torch.sigmoid(self.beta * (pos_dist - neg_dist))
|
||||
adaptive_margin = self.base_margin + self.alpha * difficulty
|
||||
|
||||
# Adaptive triplet loss using built-in ReLU
|
||||
loss = torch.relu(pos_dist - neg_dist + adaptive_margin)
|
||||
|
||||
return loss
|
||||
|
||||
batch_size = 256
|
||||
feature_dim = 1024
|
||||
|
||||
def get_inputs():
|
||||
# Generate three sets of positive vectors
|
||||
anchor = torch.abs(torch.randn(batch_size, feature_dim)) + 0.1
|
||||
positive = torch.abs(torch.randn(batch_size, feature_dim)) + 0.1
|
||||
negative = torch.abs(torch.randn(batch_size, feature_dim)) + 0.1
|
||||
return [anchor, positive, negative]
|
||||
|
||||
def get_init_inputs():
|
||||
return [1.0, 0.5, 0.1] # base_margin, alpha, beta
|
||||
|
|
@ -0,0 +1,149 @@
|
|||
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:
|
||||
python
|
||||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
… CUDA C++ source code for the kernel …
|
||||
relu_source = """""
|
||||
…
|
||||
""""
|
||||
|
||||
relu_cpp_source = """""
|
||||
torch::Tensor relu_cuda(torch::Tensor x);
|
||||
""""
|
||||
|
||||
Compile the inline CUDA code
|
||||
relu = load_inline(
|
||||
name="relu",
|
||||
cpp_sources=relu_cpp_source,
|
||||
cuda_sources=relu_source,
|
||||
functions=["relu_cuda"],
|
||||
verbose=True
|
||||
)
|
||||
|
||||
class ModelNew(torch.nn.Module):
|
||||
def init(self):
|
||||
super(ModelNew, self).init()
|
||||
self.relu = relu # The module containing the kernel
|
||||
|
||||
def forward(self, x):
|
||||
return self.relu.relu_cuda(x)
|
||||
|
||||
|
||||
You are given the following architecture:
|
||||
|
||||
python
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
Canberra Triplet Loss implementation.
|
||||
Computes triplet loss using Canberra distance as the distance metric.
|
||||
"""
|
||||
def init(self, margin=1.0):
|
||||
super(Model, self).init()
|
||||
self.margin = margin
|
||||
|
||||
def forward(self, anchor: torch.Tensor, positive: torch.Tensor, negative: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Compute Canberra Triplet Loss.
|
||||
|
||||
Args:
|
||||
anchor (torch.Tensor): Anchor samples [batch_size, feature_dim]
|
||||
positive (torch.Tensor): Positive samples [batch_size, feature_dim]
|
||||
negative (torch.Tensor): Negative samples [batch_size, feature_dim]
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Triplet loss [batch_size]
|
||||
"""
|
||||
# Input validation
|
||||
if anchor.shape != positive.shape or anchor.shape != negative.shape:
|
||||
raise ValueError(f"All input tensors must have the same shape")
|
||||
|
||||
if anchor.dim() != 2:
|
||||
raise ValueError(f"Input tensors must be 2D, got {anchor.dim()}D")
|
||||
|
||||
# Use highly optimized built-in functions
|
||||
# Combine operations to minimize memory allocation
|
||||
pos_dist = torch.sum(torch.where(
|
||||
(torch.abs(anchor) + torch.abs(positive)) > 0,
|
||||
torch.abs(anchor - positive) / (torch.abs(anchor) + torch.abs(positive)),
|
||||
torch.zeros_like(anchor)
|
||||
), dim=1)
|
||||
|
||||
neg_dist = torch.sum(torch.where(
|
||||
(torch.abs(anchor) + torch.abs(negative)) > 0,
|
||||
torch.abs(anchor - negative) / (torch.abs(anchor) + torch.abs(negative)),
|
||||
torch.zeros_like(anchor)
|
||||
), dim=1)
|
||||
|
||||
# Triplet loss using built-in ReLU
|
||||
loss = torch.relu(pos_dist - neg_dist + self.margin)
|
||||
|
||||
return loss
|
||||
|
||||
batch_size = 256
|
||||
feature_dim = 1024
|
||||
|
||||
def get_inputs():
|
||||
# Generate three sets of positive vectors
|
||||
anchor = torch.abs(torch.randn(batch_size, feature_dim)) + 0.1
|
||||
positive = torch.abs(torch.randn(batch_size, feature_dim)) + 0.1
|
||||
negative = torch.abs(torch.randn(batch_size, feature_dim)) + 0.1
|
||||
return [anchor, positive, negative]
|
||||
|
||||
def get_init_inputs():
|
||||
return [1.0] # margin
|
||||
|
||||
|
||||
|
||||
Your task is to write a new file `canberra_triplet_cudacode.py` that defines a new model `ModelNew` which uses a custom CUDA kernel to accelerate the Canberra Triplet Loss calculation. The goal is to achieve a reasonable speedup while maintaining numerical precision.
|
||||
|
||||
The recommended implementation strategy is to use a **parallel reduction pattern with operator fusion**:
|
||||
1. Launch one thread block for each sample in the batch (`batch_size` number of blocks).
|
||||
2. Within each block, have multiple threads collaborate to compute both Canberra distances (anchor-positive and anchor-negative) for that single sample in a fused manner.
|
||||
3. Each thread should iterate over the feature dimension with a stride equal to the block size, accumulating two partial sums (one for positive distance, one for negative distance) simultaneously.
|
||||
4. For each element, compute both Canberra distance terms: |anchor_i - positive_i| / (|anchor_i| + |positive_i|) and |anchor_i - negative_i| / (|anchor_i| + |negative_i|), handling the case where denominator is 0 by setting the term to 0.
|
||||
5. Use shared memory to store these partial sums (need 2 * block_size space) and then perform parallel reductions to get the final distances for that sample.
|
||||
6. The first thread of the block should compute the triplet loss: max(0, pos_dist - neg_dist + margin) using fmaxf and write the result to the output tensor.
|
||||
7. Use proper synchronization with __syncthreads() after shared memory operations.
|
||||
|
||||
Key implementation details:
|
||||
- Use extern __shared__ float sdata[] for dynamic shared memory allocation
|
||||
- Store positive partial sums in sdata[0..block_size-1] and negative partial sums in sdata[block_size..2*block_size-1]
|
||||
- Perform parallel reduction separately for both distances
|
||||
- Use TORCH_CHECK macros for input validation in the CUDA function
|
||||
- Use extra_cuda_cflags=["-O3"] for performance (avoid --use_fast_math for numerical stability)
|
||||
- Ensure the ModelNew class properly calls torch.cuda.synchronize() for accurate timing
|
||||
|
||||
The implementation should be robust, handle input validation, and focus on operator fusion to minimize memory access. The final output should be a single python file containing the CUDA kernel, the compilation logic via `load_inline`, and the new `ModelNew` class.
|
||||
|
|
@ -0,0 +1,74 @@
|
|||
###########################################################
|
||||
# 性能和精度验证程序
|
||||
###########################################################
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import time
|
||||
from canberra_adaptive_triplet_torchcode import Model, get_inputs, get_init_inputs
|
||||
from canberra_adaptive_triplet_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 canberra_adaptive_triplet 平均执行时间: {torch_time:.6f} 秒")
|
||||
print(f"自定义 CUDA canberra_adaptive_triplet 平均执行时间: {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