forked from ccf-ai-infra/GPUCodeForces
feat:add hamming_triplet #87
This commit is contained in:
parent
f989885dde
commit
149d56436b
|
|
@ -0,0 +1,113 @@
|
|||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
canberra_triplet_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
// --- 融合内核：计算每个样本的Canberra Triplet Loss ---
|
||||
__global__ void canberra_triplet_kernel_fused(
|
||||
const float* __restrict__ anchor,
|
||||
const float* __restrict__ positive,
|
||||
const float* __restrict__ negative,
|
||||
float* __restrict__ per_sample_losses, // 输出：每个样本的损失
|
||||
int batch_size,
|
||||
int feature_dim,
|
||||
float margin
|
||||
) {
|
||||
int sample_idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (sample_idx >= batch_size) return;
|
||||
|
||||
int base = sample_idx * feature_dim;
|
||||
|
||||
float pos_dist = 0.0f;
|
||||
float neg_dist = 0.0f;
|
||||
|
||||
// --- 融合计算循环：同时计算两个距离 ---
|
||||
for (int i = 0; i < feature_dim; i++) {
|
||||
float a_val = anchor[base + i];
|
||||
float p_val = positive[base + i];
|
||||
float n_val = negative[base + i];
|
||||
|
||||
// 计算正样本距离
|
||||
float pos_diff = fabsf(a_val - p_val);
|
||||
float pos_denominator = fabsf(a_val) + fabsf(p_val);
|
||||
if (pos_denominator > 1e-8f) { // 防止除以0
|
||||
pos_dist += pos_diff / pos_denominator;
|
||||
}
|
||||
|
||||
// 计算负样本距离
|
||||
float neg_diff = fabsf(a_val - n_val);
|
||||
float neg_denominator = fabsf(a_val) + fabsf(n_val);
|
||||
if (neg_denominator > 1e-8f) { // 防止除以0
|
||||
neg_dist += neg_diff / neg_denominator;
|
||||
}
|
||||
}
|
||||
|
||||
// --- 计算Triplet Loss ---
|
||||
float loss = neg_dist - pos_dist + margin;
|
||||
if (loss < 0.0f) {
|
||||
loss = 0.0f; // relu
|
||||
}
|
||||
|
||||
// 写入输出张量
|
||||
per_sample_losses[sample_idx] = loss;
|
||||
}
|
||||
|
||||
// C++包装函数：只负责启动内核，并返回包含所有样本损失的张量
|
||||
torch::Tensor canberra_triplet_cuda(torch::Tensor anchor, torch::Tensor positive, torch::Tensor negative, float margin) {
|
||||
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");
|
||||
|
||||
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::empty({batch_size}, anchor.options());
|
||||
|
||||
const int block_size = 256;
|
||||
int num_blocks = (batch_size + block_size - 1) / block_size;
|
||||
|
||||
canberra_triplet_kernel_fused<<<num_blocks, block_size>>>(
|
||||
anchor_contig.data_ptr<float>(),
|
||||
positive_contig.data_ptr<float>(),
|
||||
negative_contig.data_ptr<float>(),
|
||||
losses.data_ptr<float>(),
|
||||
batch_size,
|
||||
feature_dim,
|
||||
margin
|
||||
);
|
||||
|
||||
// 返回包含所有样本损失的张量，而不是一个标量
|
||||
return losses;
|
||||
}
|
||||
"""
|
||||
|
||||
canberra_triplet_cpp_source = """
|
||||
torch::Tensor canberra_triplet_cuda(torch::Tensor anchor, torch::Tensor positive, torch::Tensor negative, float margin);
|
||||
"""
|
||||
|
||||
canberra_triplet = load_inline(
|
||||
name="canberra_triplet_fixed",
|
||||
cpp_sources=canberra_triplet_cpp_source,
|
||||
cuda_sources=canberra_triplet_source,
|
||||
functions=["canberra_triplet_cuda"],
|
||||
extra_cuda_cflags=["-O3"],
|
||||
verbose=True
|
||||
)
|
||||
|
||||
class ModelNew(torch.nn.Module):
|
||||
def __init__(self, margin=1.0):
|
||||
super(ModelNew, self).__init__()
|
||||
self.margin = margin
|
||||
self.canberra_triplet = canberra_triplet
|
||||
|
||||
def forward(self, anchor, positive, negative):
|
||||
# CUDA内核返回每个样本的损失，我们在Python中求平均
|
||||
per_sample_losses = self.canberra_triplet.canberra_triplet_cuda(anchor, positive, negative, self.margin)
|
||||
return torch.mean(per_sample_losses)
|
||||
|
|
@ -0,0 +1,62 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
PyTorch基准实现：Canberra Triplet Loss
|
||||
"""
|
||||
def __init__(self, margin=1.0):
|
||||
super(Model, self).__init__()
|
||||
self.margin = margin
|
||||
|
||||
def forward(self, anchor, positive, negative):
|
||||
"""
|
||||
计算Canberra Triplet Loss。
|
||||
|
||||
Args:
|
||||
anchor (torch.Tensor): 锚点向量 [batch_size, feature_dim]
|
||||
positive (torch.Tensor): 正样本向量 [batch_size, feature_dim]
|
||||
negative (torch.Tensor): 负样本向量 [batch_size, feature_dim]
|
||||
|
||||
Returns:
|
||||
torch.Tensor: 平均三元组损失标量值
|
||||
"""
|
||||
# --- 第一步：计算Canberra距离 ---
|
||||
# 计算分子 |a - b|
|
||||
pos_diff = torch.abs(anchor - positive)
|
||||
neg_diff = torch.abs(anchor - negative)
|
||||
|
||||
# 计算分母 |a| + |b|
|
||||
pos_denominator = torch.abs(anchor) + torch.abs(positive)
|
||||
neg_denominator = torch.abs(anchor) + torch.abs(negative)
|
||||
|
||||
# 防止除以0
|
||||
pos_denominator = torch.clamp(pos_denominator, min=1e-8)
|
||||
neg_denominator = torch.clamp(neg_denominator, min=1e-8)
|
||||
|
||||
# 计算距离
|
||||
pos_dist = torch.sum(pos_diff / pos_denominator, dim=1)
|
||||
neg_dist = torch.sum(neg_diff / neg_denominator, dim=1)
|
||||
|
||||
# --- 第二步：计算Triplet Loss ---
|
||||
# Loss = max(0, D(a, n) - D(a, p) + margin)
|
||||
losses = torch.relu(neg_dist - pos_dist + self.margin)
|
||||
|
||||
# --- 第三步：计算平均损失 ---
|
||||
total_loss = torch.mean(losses)
|
||||
|
||||
return total_loss
|
||||
|
||||
# --- 测试数据生成函数 ---
|
||||
batch_size = 128
|
||||
feature_dim = 256
|
||||
|
||||
def get_inputs():
|
||||
# 生成非零向量以避免分母为0
|
||||
anchor = torch.randn(batch_size, feature_dim) + 0.1
|
||||
positive = torch.randn(batch_size, feature_dim) + 0.1
|
||||
negative = torch.randn(batch_size, feature_dim) + 0.1
|
||||
return [anchor, positive, negative]
|
||||
|
||||
def get_init_inputs():
|
||||
return [1.0] # margin
|
||||
|
|
@ -0,0 +1,126 @@
|
|||
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):
|
||||
"""
|
||||
Hamming Triplet Loss implementation.
|
||||
Computes the Triplet Loss using Hamming distance as the metric.
|
||||
"""
|
||||
def __init__(self, margin=1.0):
|
||||
super(Model, self).__init__()
|
||||
self.margin = margin
|
||||
|
||||
def forward(self, anchor, positive, negative):
|
||||
"""
|
||||
Compute Hamming Triplet Loss.
|
||||
|
||||
Args:
|
||||
anchor (torch.Tensor): Anchor vectors [batch_size, feature_dim]
|
||||
positive (torch.Tensor): Positive vectors [batch_size, feature_dim]
|
||||
negative (torch.Tensor): Negative vectors [batch_size, feature_dim]
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Average Triplet Loss (scalar)
|
||||
"""
|
||||
# Compute Hamming distances
|
||||
pos_dist = torch.sum(anchor != positive, dim=1).float()
|
||||
neg_dist = torch.sum(anchor != negative, dim=1).float()
|
||||
|
||||
# Compute Triplet Loss
|
||||
losses = torch.relu(neg_dist - pos_dist + self.margin)
|
||||
|
||||
# Return average loss
|
||||
return torch.mean(losses)
|
||||
|
||||
batch_size = 128
|
||||
feature_dim = 256
|
||||
|
||||
def get_inputs():
|
||||
# Generate binary vectors for Hamming distance
|
||||
anchor = torch.randint(0, 2, (batch_size, feature_dim)).float()
|
||||
positive = torch.randint(0, 2, (batch_size, feature_dim)).float()
|
||||
negative = torch.randint(0, 2, (batch_size, feature_dim)).float()
|
||||
return [anchor, positive, negative]
|
||||
|
||||
def get_init_inputs():
|
||||
return [1.0] # margin
|
||||
|
||||
|
||||
|
||||
Your task is to write a new file `hamming_triplet_cudacode.py` that defines a new model `ModelNew` which uses a custom CUDA kernel to accelerate the Hamming Triplet Loss calculation. The goal is to achieve a significant speedup while maintaining numerical precision.
|
||||
|
||||
**CRITICAL REQUIREMENTS:**
|
||||
|
||||
1. **Operator Fusion:** The entire calculation (computing two Hamming distances and the final Triplet Loss for each sample) must be performed within a **single CUDA kernel**. This kernel should output a tensor of per-sample losses.
|
||||
2. **Parallel Reduction Strategy:** The recommended implementation strategy is to use a parallel reduction pattern:
|
||||
* Launch one thread block for each sample in the batch (`batch_size` number of blocks).
|
||||
* Within each block, have multiple threads collaborate to compute the two Hamming distances for that single sample.
|
||||
* Each thread should iterate over the feature dimension with a stride equal to the block size, accumulating partial sums for both the positive and negative distances.
|
||||
* Use shared memory to store these partial sums and then perform a standard parallel reduction to get the final two distances for that sample.
|
||||
* The first thread of the block should compute the final Triplet Loss (`max(0, neg_dist - pos_dist + margin)`) and write the result to the output tensor.
|
||||
3. **Host-Side Logic:** The host-side C++ function should launch the kernel and then return the tensor of per-sample losses. The `ModelNew`'s `forward` method in Python should then call `torch.mean()` on this tensor to match the PyTorch baseline's output.
|
||||
4. **Code Structure:** The final output should be a single python file containing the CUDA kernel, the compilation logic via `load_inline`, and the new `ModelNew` class.
|
||||
5. **Compilation Flags:** Use `extra_cuda_cflags` like `"-O3"` for performance. You may also use `"--use_fast_math"` as the operations are not highly sensitive to precision.
|
||||
|
|
@ -0,0 +1,74 @@
|
|||
###########################################################
|
||||
# 性能和精度验证程序
|
||||
###########################################################
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import time
|
||||
from hamming_triplet_torchcode import Model, get_inputs, get_init_inputs
|
||||
from hamming_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 hamming_triplet 平均执行时间: {torch_time:.6f} 秒")
|
||||
print(f"自定义 CUDA hamming_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