forked from ccf-ai-infra/GPUCodeForces
FEAT:add braycurtis_adaptive_triplet #104
This commit is contained in:
parent
f989885dde
commit
eb91d5185b
|
|
@ -0,0 +1,190 @@
|
|||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
braycurtis_adaptive_triplet_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <math.h>
|
||||
|
||||
#define BLOCK_SIZE 256
|
||||
|
||||
__global__ void braycurtis_adaptive_triplet_loss_kernel(
|
||||
const float* __restrict__ anchor,
|
||||
const float* __restrict__ positive,
|
||||
const float* __restrict__ negative,
|
||||
float* __restrict__ losses,
|
||||
int batch_size,
|
||||
int feature_dim,
|
||||
float base_margin,
|
||||
float adaptive_factor,
|
||||
float min_margin,
|
||||
float max_margin
|
||||
) {
|
||||
extern __shared__ float sdata[];
|
||||
|
||||
int tid = threadIdx.x;
|
||||
int sample_idx = blockIdx.x;
|
||||
|
||||
if (sample_idx >= batch_size) return;
|
||||
|
||||
// Partial sums for both distances
|
||||
float pos_num = 0.0f, pos_den = 0.0f;
|
||||
float neg_num = 0.0f, neg_den = 0.0f;
|
||||
|
||||
int base = sample_idx * feature_dim;
|
||||
|
||||
// Each thread processes multiple elements with stride
|
||||
for (int i = tid; i < feature_dim; i += blockDim.x) {
|
||||
// Load values for better memory coalescing
|
||||
float anchor_val = anchor[base + i];
|
||||
float pos_val = positive[base + i];
|
||||
float neg_val = negative[base + i];
|
||||
|
||||
// Anchor-positive distance
|
||||
float pos_abs_diff = fabsf(anchor_val - pos_val);
|
||||
float pos_abs_sum = fabsf(anchor_val) + fabsf(pos_val);
|
||||
pos_num += pos_abs_diff;
|
||||
pos_den += pos_abs_sum;
|
||||
|
||||
// Anchor-negative distance
|
||||
float neg_abs_diff = fabsf(anchor_val - neg_val);
|
||||
float neg_abs_sum = fabsf(anchor_val) + fabsf(neg_val);
|
||||
neg_num += neg_abs_diff;
|
||||
neg_den += neg_abs_sum;
|
||||
}
|
||||
|
||||
// Store partial sums in shared memory
|
||||
float* shared_pos_num = sdata;
|
||||
float* shared_pos_den = sdata + blockDim.x;
|
||||
float* shared_neg_num = sdata + blockDim.x * 2;
|
||||
float* shared_neg_den = sdata + blockDim.x * 3;
|
||||
|
||||
shared_pos_num[tid] = pos_num;
|
||||
shared_pos_den[tid] = pos_den;
|
||||
shared_neg_num[tid] = neg_num;
|
||||
shared_neg_den[tid] = neg_den;
|
||||
|
||||
__syncthreads();
|
||||
|
||||
// Parallel reduction for all four arrays
|
||||
for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
|
||||
if (tid < stride) {
|
||||
shared_pos_num[tid] += shared_pos_num[tid + stride];
|
||||
shared_pos_den[tid] += shared_pos_den[tid + stride];
|
||||
shared_neg_num[tid] += shared_neg_num[tid + stride];
|
||||
shared_neg_den[tid] += shared_neg_den[tid + stride];
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
// First thread computes adaptive triplet loss
|
||||
if (tid == 0) {
|
||||
float pos_distance = 0.0f;
|
||||
float neg_distance = 0.0f;
|
||||
|
||||
if (shared_pos_den[0] > 0.0f) {
|
||||
pos_distance = shared_pos_num[0] / shared_pos_den[0];
|
||||
}
|
||||
|
||||
if (shared_neg_den[0] > 0.0f) {
|
||||
neg_distance = shared_neg_num[0] / shared_neg_den[0];
|
||||
}
|
||||
|
||||
// Compute adaptive margin: base_margin + adaptive_factor * (1.0 - neg_distance)
|
||||
float adaptive_margin = base_margin + adaptive_factor * (1.0f - neg_distance);
|
||||
|
||||
// Clamp margin to [min_margin, max_margin]
|
||||
if (adaptive_margin < min_margin) {
|
||||
adaptive_margin = min_margin;
|
||||
} else if (adaptive_margin > max_margin) {
|
||||
adaptive_margin = max_margin;
|
||||
}
|
||||
|
||||
// Adaptive triplet loss: max(0, pos_distance - neg_distance + adaptive_margin)
|
||||
float triplet_loss = pos_distance - neg_distance + adaptive_margin;
|
||||
losses[sample_idx] = (triplet_loss > 0.0f) ? triplet_loss : 0.0f;
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor braycurtis_adaptive_triplet_loss_cuda(
|
||||
torch::Tensor anchor,
|
||||
torch::Tensor positive,
|
||||
torch::Tensor negative,
|
||||
float base_margin,
|
||||
float adaptive_factor,
|
||||
float min_margin,
|
||||
float max_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");
|
||||
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 = block_size * 4 * sizeof(float); // For pos_num, pos_den, neg_num, neg_den
|
||||
|
||||
braycurtis_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>(),
|
||||
batch_size,
|
||||
feature_dim,
|
||||
base_margin,
|
||||
adaptive_factor,
|
||||
min_margin,
|
||||
max_margin
|
||||
);
|
||||
|
||||
return losses;
|
||||
}
|
||||
"""
|
||||
|
||||
braycurtis_adaptive_triplet_cpp_source = """
|
||||
torch::Tensor braycurtis_adaptive_triplet_loss_cuda(
|
||||
torch::Tensor anchor,
|
||||
torch::Tensor positive,
|
||||
torch::Tensor negative,
|
||||
float base_margin,
|
||||
float adaptive_factor,
|
||||
float min_margin,
|
||||
float max_margin
|
||||
);
|
||||
"""
|
||||
|
||||
braycurtis_adaptive_triplet = load_inline(
|
||||
name="braycurtis_adaptive_triplet_loss",
|
||||
cpp_sources=braycurtis_adaptive_triplet_cpp_source,
|
||||
cuda_sources=braycurtis_adaptive_triplet_source,
|
||||
functions=["braycurtis_adaptive_triplet_loss_cuda"],
|
||||
extra_cuda_cflags=["-O3"],
|
||||
verbose=True
|
||||
)
|
||||
|
||||
class ModelNew(torch.nn.Module):
|
||||
def __init__(self, base_margin=1.0, adaptive_factor=0.5, min_margin=0.1, max_margin=2.0):
|
||||
super(ModelNew, self).__init__()
|
||||
self.base_margin = base_margin
|
||||
self.adaptive_factor = adaptive_factor
|
||||
self.min_margin = min_margin
|
||||
self.max_margin = max_margin
|
||||
self.braycurtis_adaptive_triplet = braycurtis_adaptive_triplet
|
||||
|
||||
def forward(self, anchor, positive, negative):
|
||||
result = self.braycurtis_adaptive_triplet.braycurtis_adaptive_triplet_loss_cuda(
|
||||
anchor, positive, negative,
|
||||
self.base_margin, self.adaptive_factor,
|
||||
self.min_margin, self.max_margin
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
return result
|
||||
|
|
@ -0,0 +1,71 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
Bray-Curtis Adaptive Triplet Loss implementation.
|
||||
Computes triplet loss using Bray-Curtis distance with adaptive margin based on sample difficulty.
|
||||
"""
|
||||
def __init__(self, base_margin=1.0, adaptive_factor=0.5, min_margin=0.1, max_margin=2.0):
|
||||
super(Model, self).__init__()
|
||||
self.base_margin = base_margin
|
||||
self.adaptive_factor = adaptive_factor
|
||||
self.min_margin = min_margin
|
||||
self.max_margin = max_margin
|
||||
|
||||
def forward(self, anchor: torch.Tensor, positive: torch.Tensor, negative: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Compute Bray-Curtis 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 Bray-Curtis distances using highly optimized built-in functions
|
||||
# Anchor-Positive distance - fully fused computation
|
||||
pos_distance = torch.where(
|
||||
(torch.abs(anchor) + torch.abs(positive)).sum(dim=1) > 0,
|
||||
torch.abs(anchor - positive).sum(dim=1) / (torch.abs(anchor) + torch.abs(positive)).sum(dim=1),
|
||||
torch.zeros_like(anchor, device=anchor.device).sum(dim=1)
|
||||
)
|
||||
|
||||
# Anchor-Negative distance - fully fused computation
|
||||
neg_distance = torch.where(
|
||||
(torch.abs(anchor) + torch.abs(negative)).sum(dim=1) > 0,
|
||||
torch.abs(anchor - negative).sum(dim=1) / (torch.abs(anchor) + torch.abs(negative)).sum(dim=1),
|
||||
torch.zeros_like(anchor, device=anchor.device).sum(dim=1)
|
||||
)
|
||||
|
||||
# Compute adaptive margin based on negative distance (harder samples get larger margin)
|
||||
# Harder negatives (smaller neg_distance) get larger margins
|
||||
adaptive_margin = self.base_margin + self.adaptive_factor * (1.0 - neg_distance)
|
||||
adaptive_margin = torch.clamp(adaptive_margin, self.min_margin, self.max_margin)
|
||||
|
||||
# Adaptive triplet loss using built-in ReLU
|
||||
loss = torch.relu(pos_distance - neg_distance + adaptive_margin)
|
||||
|
||||
return loss
|
||||
|
||||
batch_size = 256
|
||||
feature_dim = 1024
|
||||
|
||||
def get_inputs():
|
||||
# Generate three sets of positive vectors (ecological data is typically non-negative)
|
||||
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, 2.0] # base_margin, adaptive_factor, min_margin, max_margin
|
||||
|
|
@ -0,0 +1,169 @@
|
|||
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):
|
||||
"""
|
||||
Bray-Curtis Adaptive Triplet Loss implementation.
|
||||
Computes triplet loss using Bray-Curtis distance with adaptive margin based on sample difficulty.
|
||||
"""
|
||||
def init(self, base_margin=1.0, adaptive_factor=0.5, min_margin=0.1, max_margin=2.0):
|
||||
super(Model, self).init()
|
||||
self.base_margin = base_margin
|
||||
self.adaptive_factor = adaptive_factor
|
||||
self.min_margin = min_margin
|
||||
self.max_margin = max_margin
|
||||
|
||||
def forward(self, anchor: torch.Tensor, positive: torch.Tensor, negative: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Compute Bray-Curtis 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 Bray-Curtis distances using highly optimized built-in functions
|
||||
# Anchor-Positive distance - fully fused computation
|
||||
pos_distance = torch.where(
|
||||
(torch.abs(anchor) + torch.abs(positive)).sum(dim=1) > 0,
|
||||
torch.abs(anchor - positive).sum(dim=1) / (torch.abs(anchor) + torch.abs(positive)).sum(dim=1),
|
||||
torch.zeros_like(anchor, device=anchor.device).sum(dim=1)
|
||||
)
|
||||
|
||||
# Anchor-Negative distance - fully fused computation
|
||||
neg_distance = torch.where(
|
||||
(torch.abs(anchor) + torch.abs(negative)).sum(dim=1) > 0,
|
||||
torch.abs(anchor - negative).sum(dim=1) / (torch.abs(anchor) + torch.abs(negative)).sum(dim=1),
|
||||
torch.zeros_like(anchor, device=anchor.device).sum(dim=1)
|
||||
)
|
||||
|
||||
# Compute adaptive margin based on negative distance (harder samples get larger margin)
|
||||
# Harder negatives (smaller neg_distance) get larger margins
|
||||
adaptive_margin = self.base_margin + self.adaptive_factor * (1.0 - neg_distance)
|
||||
adaptive_margin = torch.clamp(adaptive_margin, self.min_margin, self.max_margin)
|
||||
|
||||
# Adaptive triplet loss using built-in ReLU
|
||||
loss = torch.relu(pos_distance - neg_distance + adaptive_margin)
|
||||
|
||||
return loss
|
||||
|
||||
batch_size = 256
|
||||
feature_dim = 1024
|
||||
|
||||
def get_inputs():
|
||||
# Generate three sets of positive vectors (ecological data is typically non-negative)
|
||||
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, 2.0] # base_margin, adaptive_factor, min_margin, max_margin
|
||||
|
||||
|
||||
|
||||
Your task is to write a new file `braycurtis_adaptive_triplet_cudacode.py` that defines a new model `ModelNew` which uses a custom CUDA kernel to accelerate the Bray-Curtis Adaptive 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 adaptive margin computation**:
|
||||
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 Bray-Curtis distances for that single sample.
|
||||
3. Each thread should iterate over the feature dimension with a stride equal to the block size, accumulating partial sums for both distances simultaneously.
|
||||
4. For each element, compute both distance terms in parallel:
|
||||
- Anchor-Positive: |anchor_i - positive_i| / (|anchor_i| + |positive_i|)
|
||||
- Anchor-Negative: |anchor_i - negative_i| / (|anchor_i| + |negative_i|)
|
||||
5. Accumulate four partial sums: pos_num, pos_den, neg_num, neg_den
|
||||
6. Use shared memory to store these partial sums (need 4 * block_size space) and then perform parallel reductions to get the final distances for that sample.
|
||||
7. The first thread of the block should compute the adaptive margin and final loss:
|
||||
- adaptive_margin = base_margin + adaptive_factor * (1.0 - neg_distance)
|
||||
- Clamp margin to [min_margin, max_margin] using if-else statements
|
||||
- triplet_loss = max(0, pos_distance - neg_distance + adaptive_margin) using conditional expression
|
||||
8. Write the final adaptive triplet loss to the output tensor.
|
||||
|
||||
Key implementation details:
|
||||
- Use extern __shared__ float sdata[] for dynamic shared memory allocation
|
||||
- Load anchor, positive, and negative values simultaneously for better memory coalescing
|
||||
- Store partial sums in separate regions of shared memory: [pos_num, pos_den, neg_num, neg_den]
|
||||
- Perform parallel reduction separately for all four arrays in a single loop
|
||||
- Handle the adaptive margin computation within the kernel using arithmetic operations
|
||||
- Implement margin clamping using if-else statements (not built-in functions)
|
||||
- Use conditional expression (triplet_loss > 0.0f) ? triplet_loss : 0.0f for the ReLU operation
|
||||
- Use TORCH_CHECK macros for comprehensive input validation including all three input tensors
|
||||
- Use extra_cuda_cflags=["-O3"] for performance (avoid aggressive optimizations that might affect precision)
|
||||
- Use a standard block size of 256 for optimal performance
|
||||
- Ensure the ModelNew class properly calls torch.cuda.synchronize() for accurate timing
|
||||
|
||||
The implementation should be robust, handle input validation, and focus on complete operator fusion to eliminate all intermediate tensor operations. The adaptive margin computation should provide better training dynamics by adjusting margins based on sample difficulty. 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 braycurtis_adaptive_triplet_torchcode import Model, get_inputs, get_init_inputs
|
||||
from braycurtis_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 = 10000
|
||||
|
||||
# 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 braycurtis_adaptive_triplet 平均执行时间: {torch_time:.6f} 秒")
|
||||
print(f"自定义 CUDA braycurtis_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