forked from ccf-ai-infra/GPUCodeForces
finish tripletmarginloss
This commit is contained in:
parent
bee0a2a683
commit
2b8f78854c
|
|
@ -0,0 +1,79 @@
|
|||
The standard PyTorch implementation of Triplet Margin Loss, torch.nn.TripletMarginLoss, is a textbook example of a memory-bandwidth-bound operation. Its calculation involves a long chain of simple, element-wise operations that generate numerous large intermediate tensors, leading to severe performance degradation.
|
||||
|
||||
The original process for a batch of triplets can be broken down as follows:
|
||||
|
||||
Positive Pair Distance:
|
||||
|
||||
d_pos = anchor - positive (Creates intermediate tensor)
|
||||
|
||||
d_pos = d_pos.pow(2) (Creates intermediate tensor)
|
||||
|
||||
d_pos = d_pos.sum(dim=1) (Reduction, creates intermediate tensor)
|
||||
|
||||
d_pos = d_pos.sqrt() (Creates intermediate tensor)
|
||||
|
||||
Negative Pair Distance:
|
||||
|
||||
d_neg = anchor - negative (Creates intermediate tensor)
|
||||
|
||||
d_neg = d_neg.pow(2) (Creates intermediate tensor)
|
||||
|
||||
d_neg = d_neg.sum(dim=1) (Reduction, creates intermediate tensor)
|
||||
|
||||
d_neg = d_neg.sqrt() (Creates intermediate tensor)
|
||||
|
||||
Final Loss Calculation:
|
||||
|
||||
loss = d_pos - d_neg + margin (Creates intermediate tensor)
|
||||
|
||||
loss = torch.clamp(loss, min=0) (Final operation)
|
||||
|
||||
This sequence triggers at least 10 separate CUDA kernel launches and forces the GPU to write and read gigabytes of temporary data to and from its global memory, while the actual arithmetic computation is minimal.
|
||||
|
||||
You should fuse this entire computational graph into a single, highly efficient CUDA kernel. The kernel will compute the loss for each triplet in a single pass, keeping all intermediate values within high-speed registers and shared memory, thus eliminating the global memory bottleneck.
|
||||
|
||||
Considerations:
|
||||
|
||||
Parallelization Strategy: The ideal approach is to assign one CUDA block to compute the loss for one triplet in the batch. The grid dimension will be equal to the batch size.
|
||||
|
||||
Parallel Reduction for L2 Distance: The sum() operation within the L2 distance calculation is a classic parallel reduction problem. You must implement an efficient reduction inside the CUDA kernel using shared memory.
|
||||
|
||||
Threads within a block will collaboratively load slices of the anchor, positive, and negative vectors into shared memory.
|
||||
|
||||
Each thread computes the squared difference for a subset of the feature dimension.
|
||||
|
||||
A synchronized, tree-based reduction is performed within the block using the shared memory array to sum up all the partial results.
|
||||
|
||||
Final Calculation in Registers: Once the positive and negative distances are calculated via reduction (the results will likely reside in thread 0's registers), that same thread will perform the final max(0, d_pos - d_neg + margin) calculation before writing the single scalar result back to global memory.
|
||||
|
||||
You are given the following baseline architecture:
|
||||
|
||||
```python
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, margin: float = 1.0):
|
||||
super(Model, self).__init__()
|
||||
self.margin = margin
|
||||
self.triplet_margin_loss = torch.nn.TripletMarginLoss(margin=self.margin, reduction='mean')
|
||||
|
||||
def forward(self, anchor: torch.Tensor, positive: torch.Tensor, negative: torch.Tensor) -> torch.Tensor:
|
||||
return self.triplet_margin_loss(anchor, positive, negative)
|
||||
|
||||
batch_size = 512
|
||||
dim = 4096
|
||||
margin = 1.0
|
||||
|
||||
def get_inputs():
|
||||
"""
|
||||
为anchor, positive, 和 negative生成三个随机张量。
|
||||
"""
|
||||
anchor = torch.randn(batch_size, dim)
|
||||
positive = torch.randn(batch_size, dim)
|
||||
negative = torch.randn(batch_size, dim)
|
||||
return [anchor, positive, negative]
|
||||
|
||||
def get_init_inputs():
|
||||
return [margin]
|
||||
|
|
@ -0,0 +1,88 @@
|
|||
###########################################################
|
||||
# 性能和精度验证程序
|
||||
###########################################################
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import time
|
||||
from tripletmarginloss_torch import Model, get_inputs, get_init_inputs
|
||||
from tripletmarginloss_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)
|
||||
|
||||
# 更严格的精度检查
|
||||
abs_diff = (output_torch - output_cuda).abs()
|
||||
max_diff = abs_diff.max().item()
|
||||
mean_diff = abs_diff.mean().item()
|
||||
|
||||
print(f"最大差异: {max_diff:.6f}")
|
||||
print(f"平均差异: {mean_diff:.6f}")
|
||||
|
||||
precision_flag = torch.allclose(output_torch, output_cuda, rtol=1e-05, atol=1e-05)
|
||||
if precision_flag:
|
||||
print("✅ 精度对齐：两个模型的输出结果非常接近。")
|
||||
else:
|
||||
print("❌ 精度不一致！")
|
||||
|
||||
print("\n-------------------- 性能加速比测试 --------------------")
|
||||
num_iterations = 1000 # 增加迭代次数以获得更准确的时间测量
|
||||
|
||||
# Warm up
|
||||
for _ in range(100):
|
||||
_ = torch_model(*inputs)
|
||||
_ = cuda_model(*inputs)
|
||||
|
||||
# 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 (matmul + relu) 平均执行时间: {torch_time:.6f} 秒")
|
||||
print(f"自定义 CUDA ReLU 平均执行时间: {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()
|
||||
|
|
@ -0,0 +1,108 @@
|
|||
# example_cudacode.py
|
||||
|
||||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
# Triplet Margin Loss的自定义CUDA实现
|
||||
triplet_loss_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <cmath>
|
||||
|
||||
// 定义共享内存和块的大小
|
||||
#define BLOCK_SIZE 256
|
||||
|
||||
// 计算两个向量之间的L2距离的平方
|
||||
__device__ float squared_l2_distance(const float* v1, const float* v2, int dim, float* sdata) {
|
||||
float my_sum = 0.0f;
|
||||
// 每个线程计算一部分元素的平方差之和
|
||||
for (int i = threadIdx.x; i < dim; i += blockDim.x) {
|
||||
float diff = v1[i] - v2[i];
|
||||
my_sum += diff * diff;
|
||||
}
|
||||
sdata[threadIdx.x] = my_sum;
|
||||
__syncthreads(); // 确保所有线程都完成了它们的初始求和
|
||||
|
||||
// 使用共享内存执行并行规约
|
||||
for (unsigned int s = blockDim.x / 2; s > 0; s >>= 1) {
|
||||
if (threadIdx.x < s) {
|
||||
sdata[threadIdx.x] += sdata[threadIdx.x + s];
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
// 最终的平方和在sdata[0]中
|
||||
return sdata[0];
|
||||
}
|
||||
|
||||
__global__ void triplet_loss_kernel(
|
||||
const float* anchor, const float* positive, const float* negative,
|
||||
float* loss, int batch_size, int dim, float margin) {
|
||||
|
||||
int batch_idx = blockIdx.x;
|
||||
if (batch_idx >= batch_size) return;
|
||||
|
||||
__shared__ float sdata[BLOCK_SIZE];
|
||||
|
||||
// 计算当前三元组的向量指针
|
||||
const float* anchor_ptr = anchor + batch_idx * dim;
|
||||
const float* positive_ptr = positive + batch_idx * dim;
|
||||
const float* negative_ptr = negative + batch_idx * dim;
|
||||
|
||||
// 计算正对和负对的距离
|
||||
float dist_pos_sq = squared_l2_distance(anchor_ptr, positive_ptr, dim, sdata);
|
||||
float dist_neg_sq = squared_l2_distance(anchor_ptr, negative_ptr, dim, sdata);
|
||||
|
||||
// 只有块中的第一个线程执行最终的计算和写入操作
|
||||
if (threadIdx.x == 0) {
|
||||
float dist_pos = sqrtf(dist_pos_sq);
|
||||
float dist_neg = sqrtf(dist_neg_sq);
|
||||
float loss_val = dist_pos - dist_neg + margin;
|
||||
|
||||
loss[batch_idx] = fmaxf(0.0f, loss_val);
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor triplet_loss_cuda(
|
||||
torch::Tensor anchor, torch::Tensor positive, torch::Tensor negative, float margin) {
|
||||
|
||||
int batch_size = anchor.size(0);
|
||||
int dim = anchor.size(1);
|
||||
|
||||
auto options = torch::TensorOptions().device(anchor.device()).dtype(anchor.dtype());
|
||||
auto loss = torch::empty({batch_size}, options);
|
||||
|
||||
dim3 grid(batch_size);
|
||||
|
||||
dim3 block(BLOCK_SIZE);
|
||||
|
||||
triplet_loss_kernel<<<grid, block>>>(
|
||||
anchor.data_ptr<float>(), positive.data_ptr<float>(), negative.data_ptr<float>(),
|
||||
loss.data_ptr<float>(), batch_size, dim, margin);
|
||||
|
||||
return loss.sum()/batch_size;
|
||||
}
|
||||
"""
|
||||
|
||||
triplet_loss_cpp_source = """
|
||||
torch::Tensor triplet_loss_cuda(
|
||||
torch::Tensor anchor, torch::Tensor positive, torch::Tensor negative, float margin);
|
||||
"""
|
||||
|
||||
# 编译内联CUDA代码
|
||||
triplet_loss_module = load_inline(
|
||||
name="triplet_loss_module",
|
||||
cpp_sources=triplet_loss_cpp_source,
|
||||
cuda_sources=triplet_loss_source,
|
||||
functions=["triplet_loss_cuda"],
|
||||
verbose=True
|
||||
)
|
||||
|
||||
class ModelNew(torch.nn.Module):
|
||||
def __init__(self, margin: float = 1.0):
|
||||
super(ModelNew, self).__init__()
|
||||
self.margin = margin
|
||||
self.triplet_loss = triplet_loss_module
|
||||
|
||||
def forward(self, anchor, positive, negative):
|
||||
return self.triplet_loss.triplet_loss_cuda(anchor, positive, negative, self.margin)
|
||||
|
|
@ -0,0 +1,41 @@
|
|||
# example_torchcode.py
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
一个计算Triplet Margin Loss的简单模型。
|
||||
"""
|
||||
def __init__(self, margin: float = 1.0):
|
||||
super(Model, self).__init__()
|
||||
self.margin = margin
|
||||
self.triplet_margin_loss = torch.nn.TripletMarginLoss(margin=self.margin, reduction='mean')
|
||||
|
||||
def forward(self, anchor: torch.Tensor, positive: torch.Tensor, negative: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
计算三元组损失。
|
||||
使用 reduction='none' 来为批次中的每个样本生成一个损失值，以便与CUDA内核进行比较。
|
||||
"""
|
||||
return self.triplet_margin_loss(anchor, positive, negative)
|
||||
|
||||
# 定义标准维度
|
||||
batch_size = 512
|
||||
dim = 4096
|
||||
margin = 1.0
|
||||
|
||||
def get_inputs():
|
||||
"""
|
||||
为anchor, positive, 和 negative生成三个随机张量。
|
||||
"""
|
||||
anchor = torch.randn(batch_size, dim)
|
||||
positive = torch.randn(batch_size, dim)
|
||||
negative = torch.randn(batch_size, dim)
|
||||
return [anchor, positive, negative]
|
||||
|
||||
def get_init_inputs():
|
||||
"""
|
||||
提供模型初始化所需的margin。
|
||||
"""
|
||||
return [margin]
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
|
|
@ -0,0 +1,79 @@
|
|||
The standard PyTorch implementation of Triplet Margin Loss, torch.nn.TripletMarginLoss, is a textbook example of a memory-bandwidth-bound operation. Its calculation involves a long chain of simple, element-wise operations that generate numerous large intermediate tensors, leading to severe performance degradation.
|
||||
|
||||
The original process for a batch of triplets can be broken down as follows:
|
||||
|
||||
Positive Pair Distance:
|
||||
|
||||
d_pos = anchor - positive (Creates intermediate tensor)
|
||||
|
||||
d_pos = d_pos.pow(2) (Creates intermediate tensor)
|
||||
|
||||
d_pos = d_pos.sum(dim=1) (Reduction, creates intermediate tensor)
|
||||
|
||||
d_pos = d_pos.sqrt() (Creates intermediate tensor)
|
||||
|
||||
Negative Pair Distance:
|
||||
|
||||
d_neg = anchor - negative (Creates intermediate tensor)
|
||||
|
||||
d_neg = d_neg.pow(2) (Creates intermediate tensor)
|
||||
|
||||
d_neg = d_neg.sum(dim=1) (Reduction, creates intermediate tensor)
|
||||
|
||||
d_neg = d_neg.sqrt() (Creates intermediate tensor)
|
||||
|
||||
Final Loss Calculation:
|
||||
|
||||
loss = d_pos - d_neg + margin (Creates intermediate tensor)
|
||||
|
||||
loss = torch.clamp(loss, min=0) (Final operation)
|
||||
|
||||
This sequence triggers at least 10 separate CUDA kernel launches and forces the GPU to write and read gigabytes of temporary data to and from its global memory, while the actual arithmetic computation is minimal.
|
||||
|
||||
You should fuse this entire computational graph into a single, highly efficient CUDA kernel. The kernel will compute the loss for each triplet in a single pass, keeping all intermediate values within high-speed registers and shared memory, thus eliminating the global memory bottleneck.
|
||||
|
||||
Considerations:
|
||||
|
||||
Parallelization Strategy: The ideal approach is to assign one CUDA block to compute the loss for one triplet in the batch. The grid dimension will be equal to the batch size.
|
||||
|
||||
Parallel Reduction for L2 Distance: The sum() operation within the L2 distance calculation is a classic parallel reduction problem. You must implement an efficient reduction inside the CUDA kernel using shared memory.
|
||||
|
||||
Threads within a block will collaboratively load slices of the anchor, positive, and negative vectors into shared memory.
|
||||
|
||||
Each thread computes the squared difference for a subset of the feature dimension.
|
||||
|
||||
A synchronized, tree-based reduction is performed within the block using the shared memory array to sum up all the partial results.
|
||||
|
||||
Final Calculation in Registers: Once the positive and negative distances are calculated via reduction (the results will likely reside in thread 0's registers), that same thread will perform the final max(0, d_pos - d_neg + margin) calculation before writing the single scalar result back to global memory.
|
||||
|
||||
You are given the following baseline architecture:
|
||||
|
||||
```python
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, margin: float = 1.0):
|
||||
super(Model, self).__init__()
|
||||
self.margin = margin
|
||||
self.triplet_margin_loss = torch.nn.TripletMarginLoss(margin=self.margin, reduction='mean')
|
||||
|
||||
def forward(self, anchor: torch.Tensor, positive: torch.Tensor, negative: torch.Tensor) -> torch.Tensor:
|
||||
return self.triplet_margin_loss(anchor, positive, negative)
|
||||
|
||||
batch_size = 512
|
||||
dim = 4096
|
||||
margin = 1.0
|
||||
|
||||
def get_inputs():
|
||||
"""
|
||||
为anchor, positive, 和 negative生成三个随机张量。
|
||||
"""
|
||||
anchor = torch.randn(batch_size, dim)
|
||||
positive = torch.randn(batch_size, dim)
|
||||
negative = torch.randn(batch_size, dim)
|
||||
return [anchor, positive, negative]
|
||||
|
||||
def get_init_inputs():
|
||||
return [margin]
|
||||
|
|
@ -0,0 +1,88 @@
|
|||
###########################################################
|
||||
# 性能和精度验证程序
|
||||
###########################################################
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import time
|
||||
from tripletmarginloss_torch import Model, get_inputs, get_init_inputs
|
||||
from tripletmarginloss_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)
|
||||
|
||||
# 更严格的精度检查
|
||||
abs_diff = (output_torch - output_cuda).abs()
|
||||
max_diff = abs_diff.max().item()
|
||||
mean_diff = abs_diff.mean().item()
|
||||
|
||||
print(f"最大差异: {max_diff:.6f}")
|
||||
print(f"平均差异: {mean_diff:.6f}")
|
||||
|
||||
precision_flag = torch.allclose(output_torch, output_cuda, rtol=1e-05, atol=1e-05)
|
||||
if precision_flag:
|
||||
print("✅ 精度对齐：两个模型的输出结果非常接近。")
|
||||
else:
|
||||
print("❌ 精度不一致！")
|
||||
|
||||
print("\n-------------------- 性能加速比测试 --------------------")
|
||||
num_iterations = 1000 # 增加迭代次数以获得更准确的时间测量
|
||||
|
||||
# Warm up
|
||||
for _ in range(100):
|
||||
_ = torch_model(*inputs)
|
||||
_ = cuda_model(*inputs)
|
||||
|
||||
# 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 (matmul + relu) 平均执行时间: {torch_time:.6f} 秒")
|
||||
print(f"自定义 CUDA ReLU 平均执行时间: {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()
|
||||
|
|
@ -0,0 +1,108 @@
|
|||
# example_cudacode.py
|
||||
|
||||
import torch
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
# Triplet Margin Loss的自定义CUDA实现
|
||||
triplet_loss_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <cmath>
|
||||
|
||||
// 定义共享内存和块的大小
|
||||
#define BLOCK_SIZE 256
|
||||
|
||||
// 计算两个向量之间的L2距离的平方
|
||||
__device__ float squared_l2_distance(const float* v1, const float* v2, int dim, float* sdata) {
|
||||
float my_sum = 0.0f;
|
||||
// 每个线程计算一部分元素的平方差之和
|
||||
for (int i = threadIdx.x; i < dim; i += blockDim.x) {
|
||||
float diff = v1[i] - v2[i];
|
||||
my_sum += diff * diff;
|
||||
}
|
||||
sdata[threadIdx.x] = my_sum;
|
||||
__syncthreads(); // 确保所有线程都完成了它们的初始求和
|
||||
|
||||
// 使用共享内存执行并行规约
|
||||
for (unsigned int s = blockDim.x / 2; s > 0; s >>= 1) {
|
||||
if (threadIdx.x < s) {
|
||||
sdata[threadIdx.x] += sdata[threadIdx.x + s];
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
// 最终的平方和在sdata[0]中
|
||||
return sdata[0];
|
||||
}
|
||||
|
||||
__global__ void triplet_loss_kernel(
|
||||
const float* anchor, const float* positive, const float* negative,
|
||||
float* loss, int batch_size, int dim, float margin) {
|
||||
|
||||
int batch_idx = blockIdx.x;
|
||||
if (batch_idx >= batch_size) return;
|
||||
|
||||
__shared__ float sdata[BLOCK_SIZE];
|
||||
|
||||
// 计算当前三元组的向量指针
|
||||
const float* anchor_ptr = anchor + batch_idx * dim;
|
||||
const float* positive_ptr = positive + batch_idx * dim;
|
||||
const float* negative_ptr = negative + batch_idx * dim;
|
||||
|
||||
// 计算正对和负对的距离
|
||||
float dist_pos_sq = squared_l2_distance(anchor_ptr, positive_ptr, dim, sdata);
|
||||
float dist_neg_sq = squared_l2_distance(anchor_ptr, negative_ptr, dim, sdata);
|
||||
|
||||
// 只有块中的第一个线程执行最终的计算和写入操作
|
||||
if (threadIdx.x == 0) {
|
||||
float dist_pos = sqrtf(dist_pos_sq);
|
||||
float dist_neg = sqrtf(dist_neg_sq);
|
||||
float loss_val = dist_pos - dist_neg + margin;
|
||||
|
||||
loss[batch_idx] = fmaxf(0.0f, loss_val);
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor triplet_loss_cuda(
|
||||
torch::Tensor anchor, torch::Tensor positive, torch::Tensor negative, float margin) {
|
||||
|
||||
int batch_size = anchor.size(0);
|
||||
int dim = anchor.size(1);
|
||||
|
||||
auto options = torch::TensorOptions().device(anchor.device()).dtype(anchor.dtype());
|
||||
auto loss = torch::empty({batch_size}, options);
|
||||
|
||||
dim3 grid(batch_size);
|
||||
|
||||
dim3 block(BLOCK_SIZE);
|
||||
|
||||
triplet_loss_kernel<<<grid, block>>>(
|
||||
anchor.data_ptr<float>(), positive.data_ptr<float>(), negative.data_ptr<float>(),
|
||||
loss.data_ptr<float>(), batch_size, dim, margin);
|
||||
|
||||
return loss.sum()/batch_size;
|
||||
}
|
||||
"""
|
||||
|
||||
triplet_loss_cpp_source = """
|
||||
torch::Tensor triplet_loss_cuda(
|
||||
torch::Tensor anchor, torch::Tensor positive, torch::Tensor negative, float margin);
|
||||
"""
|
||||
|
||||
# 编译内联CUDA代码
|
||||
triplet_loss_module = load_inline(
|
||||
name="triplet_loss_module",
|
||||
cpp_sources=triplet_loss_cpp_source,
|
||||
cuda_sources=triplet_loss_source,
|
||||
functions=["triplet_loss_cuda"],
|
||||
verbose=True
|
||||
)
|
||||
|
||||
class ModelNew(torch.nn.Module):
|
||||
def __init__(self, margin: float = 1.0):
|
||||
super(ModelNew, self).__init__()
|
||||
self.margin = margin
|
||||
self.triplet_loss = triplet_loss_module
|
||||
|
||||
def forward(self, anchor, positive, negative):
|
||||
return self.triplet_loss.triplet_loss_cuda(anchor, positive, negative, self.margin)
|
||||
|
|
@ -0,0 +1,29 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, margin: float = 1.0):
|
||||
super(Model, self).__init__()
|
||||
self.margin = margin
|
||||
self.triplet_margin_loss = torch.nn.TripletMarginLoss(margin=self.margin, reduction='mean')
|
||||
|
||||
def forward(self, anchor: torch.Tensor, positive: torch.Tensor, negative: torch.Tensor) -> torch.Tensor:
|
||||
return self.triplet_margin_loss(anchor, positive, negative)
|
||||
|
||||
# 定义标准维度
|
||||
batch_size = 512
|
||||
dim = 4096
|
||||
margin = 1.0
|
||||
|
||||
def get_inputs():
|
||||
"""
|
||||
为anchor, positive, 和 negative生成三个随机张量。
|
||||
"""
|
||||
anchor = torch.randn(batch_size, dim)
|
||||
positive = torch.randn(batch_size, dim)
|
||||
negative = torch.randn(batch_size, dim)
|
||||
return [anchor, positive, negative]
|
||||
|
||||
def get_init_inputs():
|
||||
return [margin]
|
||||
Loading…
Reference in New Issue