forked from ccf-ai-infra/GPUCodeForces
Merge pull request 'finish multilabelmarginloss #44' (#73) from hli28146/GPUCodeForces:MultiLabelMarginLoss into main
This commit is contained in:
commit
3a164824d6
|
|
@ -0,0 +1,156 @@
|
|||
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 multi_label_margin_loss_cuda_forward(
|
||||
const torch::Tensor& input,
|
||||
const torch::Tensor& target,
|
||||
const std::string& reduction
|
||||
);
|
||||
"""
|
||||
|
||||
cuda_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#define BLOCK_SIZE 256
|
||||
// 设置一个合理的单个样本最大正类标签数，用于共享内存数组
|
||||
#define MAX_POSITIVE_LABELS 32
|
||||
|
||||
template <typename T>
|
||||
__global__ void multi_label_margin_loss_kernel(
|
||||
T* output, // (N)
|
||||
const T* input, // (N, C)
|
||||
const long* target, // (N, C)
|
||||
const int N,
|
||||
const int C)
|
||||
{
|
||||
int sample_idx = blockIdx.x;
|
||||
if (sample_idx >= N) return;
|
||||
|
||||
__shared__ long positive_indices[MAX_POSITIVE_LABELS];
|
||||
__shared__ int num_positives;
|
||||
|
||||
// 线程 0 初始化共享内存计数器
|
||||
if (threadIdx.x == 0) {
|
||||
num_positives = 0;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (int i = threadIdx.x; i < C; i += blockDim.x) {
|
||||
long label = target[sample_idx * C + i];
|
||||
if (label != -1) {
|
||||
int index = atomicAdd(&num_positives, 1);
|
||||
if (index < MAX_POSITIVE_LABELS) {
|
||||
positive_indices[index] = label;
|
||||
}
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
__shared__ T sdata[BLOCK_SIZE];
|
||||
int tid = threadIdx.x;
|
||||
T my_sum = 0.0f;
|
||||
const T* input_row = input + sample_idx * C;
|
||||
|
||||
for (int neg_class_idx = tid; neg_class_idx < C; neg_class_idx += blockDim.x) {
|
||||
// 检查当前类别是否为正类
|
||||
bool is_positive = false;
|
||||
for (int j = 0; j < num_positives; ++j) {
|
||||
if (positive_indices[j] == neg_class_idx) {
|
||||
is_positive = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// 如果是负类，则计算与所有正类的损失
|
||||
if (!is_positive) {
|
||||
T x_neg = input_row[neg_class_idx];
|
||||
for (int j = 0; j < num_positives; ++j) {
|
||||
long pos_class_idx = positive_indices[j];
|
||||
T x_pos = input_row[pos_class_idx];
|
||||
T loss_term = 1.0f - (x_pos - x_neg);
|
||||
if (loss_term > 0) {
|
||||
my_sum += loss_term;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
sdata[tid] = my_sum;
|
||||
__syncthreads();
|
||||
|
||||
for (int s = blockDim.x / 2; s > 0; s >>= 1) {
|
||||
if (tid < s) {
|
||||
sdata[tid] += sdata[tid + s];
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
if (tid == 0) {
|
||||
output[sample_idx] = sdata[0] / C;
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor multi_label_margin_loss_cuda_forward(
|
||||
const torch::Tensor& input,
|
||||
const torch::Tensor& target,
|
||||
const std::string& reduction)
|
||||
{
|
||||
TORCH_CHECK(input.is_cuda() && target.is_cuda(), "Tensors must be on CUDA");
|
||||
TORCH_CHECK(input.dim() == 2, "Input must be 2D");
|
||||
TORCH_CHECK(target.dim() == 2, "Target must be 2D");
|
||||
TORCH_CHECK(input.size(0) == target.size(0), "Batch sizes must match");
|
||||
TORCH_CHECK(input.is_contiguous() && target.is_contiguous(), "Tensors must be contiguous");
|
||||
|
||||
const int N = input.size(0);
|
||||
const int C = input.size(1);
|
||||
|
||||
auto options = torch::TensorOptions().device(input.device()).dtype(input.dtype());
|
||||
auto sample_losses = torch::empty({N}, options);
|
||||
|
||||
dim3 grid(N);
|
||||
dim3 block(BLOCK_SIZE);
|
||||
|
||||
AT_DISPATCH_FLOATING_TYPES(input.scalar_type(), "multi_label_margin_loss_kernel", ([&] {{
|
||||
multi_label_margin_loss_kernel<scalar_t><<<grid, block>>>(
|
||||
sample_losses.data_ptr<scalar_t>(),
|
||||
input.data_ptr<scalar_t>(),
|
||||
target.data_ptr<long>(),
|
||||
N, C
|
||||
);
|
||||
}}));
|
||||
|
||||
if (reduction == "none") {
|
||||
return sample_losses;
|
||||
} else if (reduction == "sum") {
|
||||
return sample_losses.sum();
|
||||
} else {{ // "mean"
|
||||
return sample_losses.mean();
|
||||
}}
|
||||
}
|
||||
"""
|
||||
|
||||
class ModelNew(nn.Module):
|
||||
def __init__(self, reduction='mean'):
|
||||
super(ModelNew, self).__init__()
|
||||
self.reduction = reduction
|
||||
|
||||
self.op = load_inline(
|
||||
name='multi_label_margin_loss_op',
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=cuda_source,
|
||||
functions=['multi_label_margin_loss_cuda_forward'],
|
||||
verbose=False
|
||||
)
|
||||
|
||||
def forward(self, input_tensor: torch.Tensor, target_tensor: torch.Tensor) -> torch.Tensor:
|
||||
return self.op.multi_label_margin_loss_cuda_forward(
|
||||
input_tensor,
|
||||
target_tensor,
|
||||
self.reduction
|
||||
)
|
||||
|
|
@ -0,0 +1,34 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
import numpy as np
|
||||
|
||||
BATCH_SIZE = 512
|
||||
NUM_CLASSES = 1024
|
||||
REDUCTION = 'mean'
|
||||
# 每个样本的正类标签数量范围
|
||||
MIN_LABELS = 1
|
||||
MAX_LABELS = 10
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, reduction='mean'):
|
||||
super(Model, self).__init__()
|
||||
self.loss_fn = nn.MultiLabelMarginLoss(reduction=reduction)
|
||||
|
||||
def forward(self, input_tensor: torch.Tensor, target_tensor: torch.Tensor) -> torch.Tensor:
|
||||
return self.loss_fn(input_tensor, target_tensor)
|
||||
|
||||
def get_inputs():
|
||||
input_tensor = torch.randn(BATCH_SIZE, NUM_CLASSES, dtype=torch.float32)
|
||||
|
||||
# target每行包含正类索引，并用 -1 填充
|
||||
target_np = np.full((BATCH_SIZE, NUM_CLASSES), -1, dtype=np.int64)
|
||||
for i in range(BATCH_SIZE):
|
||||
num_labels = np.random.randint(MIN_LABELS, MAX_LABELS + 1)
|
||||
labels = np.random.choice(NUM_CLASSES, num_labels, replace=False)
|
||||
target_np[i, :num_labels] = labels
|
||||
target_tensor = torch.from_numpy(target_np)
|
||||
|
||||
return [input_tensor.contiguous(), target_tensor.contiguous()]
|
||||
|
||||
def get_init_inputs():
|
||||
return [REDUCTION]
|
||||
|
|
@ -0,0 +1,66 @@
|
|||
Write a custom CUDA kernel to optimize `torch.nn.MultiLabelMarginLoss`.
|
||||
|
||||
The original operation is defined by the formula:
|
||||
`loss(x, y) = sum_{j,i} max(0, 1 - (x[y[j]] - x[i])) / (x.size(0) * x.size(1))`
|
||||
where `y[j]` are the positive class indices for a sample and `i` are the negative class indices. This is computed per sample and then reduced.
|
||||
|
||||
**Problem Analysis:**
|
||||
`MultiLabelMarginLoss` is notoriously difficult to vectorize efficiently in PyTorch. The performance bottlenecks are severe:
|
||||
1. **Irregular Data Access**: Each sample has a variable number of positive labels defined in `y`, which are padded with -1. Identifying the set of positive and negative classes for each sample requires complex, non-vectorized logic (e.g., masks, loops, or boolean indexing), which is slow.
|
||||
2. **Massive Intermediate Tensors**: A naive vectorized approach would require gathering scores for positive classes and broadcasting them for subtraction against scores of negative classes. This would create huge intermediate tensors and is highly memory-inefficient.
|
||||
3. **Complex Nested Loop Logic**: The core formula is a nested loop (`for each positive class`, `for each negative class`) for every sample, which is antithetical to efficient GPU execution without a custom kernel.
|
||||
|
||||
**Optimization Strategy: Fused Block-Level Parallelism with Shared Memory Caching**
|
||||
|
||||
The strategy is to implement the entire complex logic within a single CUDA kernel, using a block-per-sample parallelization model.
|
||||
|
||||
1. **Parallelization Model**: A grid of `N` blocks is launched, where `N` is the batch size. Each thread block is assigned to compute the total loss for one sample.
|
||||
|
||||
2. **Shared Memory Caching**: For each sample (block), the kernel first collaboratively reads the list of positive class indices from the `target` tensor. These indices (and their count) are cached in **shared memory**. This makes the critical metadata for the sample instantly accessible to all threads in the block.
|
||||
|
||||
3. **Fused Computation Loop**: The threads within a block then work together to iterate through all `C` possible classes. For each class `i`, a thread checks if it's a positive or negative class using the cached shared memory data.
|
||||
* If `i` is a negative class, the thread then iterates through the *positive class indices cached in shared memory*.
|
||||
* For each positive-negative pair, it calculates the hinge loss term `max(0, 1 - (x_pos - x_neg))` and accumulates it into a thread-local register. This fuses the nested loops, indexing, subtraction, and `max` operations.
|
||||
|
||||
4. **Efficient Intra-Block Reduction**: Once all classes are processed, a fast parallel reduction is performed using shared memory to sum the partial results from all threads within the block into a single total loss for that sample.
|
||||
|
||||
5. **Finalization**: The first thread of each block performs the final division and writes the result to the output tensor. The kernel directly produces the per-sample losses (`reduction='none'`). The final batch reduction (`'mean'` or `'sum'`) is efficiently handled by a single PyTorch call on the small 1D output tensor.
|
||||
|
||||
This approach transforms the complex, memory-bound, and hard-to-vectorize PyTorch operation into a single, efficient, compute-focused CUDA kernel.
|
||||
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 numpy as np
|
||||
|
||||
BATCH_SIZE = 512
|
||||
NUM_CLASSES = 1024
|
||||
REDUCTION = 'mean'
|
||||
# 每个样本的正类标签数量范围
|
||||
MIN_LABELS = 1
|
||||
MAX_LABELS = 10
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, reduction='mean'):
|
||||
super(Model, self).__init__()
|
||||
self.loss_fn = nn.MultiLabelMarginLoss(reduction=reduction)
|
||||
|
||||
def forward(self, input_tensor: torch.Tensor, target_tensor: torch.Tensor) -> torch.Tensor:
|
||||
return self.loss_fn(input_tensor, target_tensor)
|
||||
|
||||
def get_inputs():
|
||||
input_tensor = torch.randn(BATCH_SIZE, NUM_CLASSES, dtype=torch.float32)
|
||||
|
||||
# target每行包含正类索引，并用 -1 填充
|
||||
target_np = np.full((BATCH_SIZE, NUM_CLASSES), -1, dtype=np.int64)
|
||||
for i in range(BATCH_SIZE):
|
||||
num_labels = np.random.randint(MIN_LABELS, MAX_LABELS + 1)
|
||||
labels = np.random.choice(NUM_CLASSES, num_labels, replace=False)
|
||||
target_np[i, :num_labels] = labels
|
||||
target_tensor = torch.from_numpy(target_np)
|
||||
|
||||
return [input_tensor.contiguous(), target_tensor.contiguous()]
|
||||
|
||||
def get_init_inputs():
|
||||
return [REDUCTION]
|
||||
|
|
@ -0,0 +1,88 @@
|
|||
###########################################################
|
||||
# 性能和精度验证程序
|
||||
###########################################################
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import time
|
||||
from multilabelmarginloss_torch import Model, get_inputs, get_init_inputs
|
||||
from multilabelmarginloss_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 MultiLabelMarginLoss 平均执行时间: {torch_time:.6f} 秒")
|
||||
print(f"自定义 CUDA MultiLabelMarginLoss 平均执行时间: {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