forked from ccf-ai-infra/GPUCodeForces
167 lines
4.0 KiB
Python
167 lines
4.0 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from torch.utils.cpp_extension import load_inline
|
|
|
|
cuda_source = """
|
|
#include <torch/extension.h>
|
|
#include <cuda_runtime.h>
|
|
#include <math.h>
|
|
|
|
__global__ void fast_ap_loss_kernel(
|
|
const float* __restrict__ x,
|
|
const long* __restrict__ labels,
|
|
float* __restrict__ ap_out,
|
|
int batch_size,
|
|
int dim,
|
|
int num_bins,
|
|
float max_dist)
|
|
{
|
|
int i = blockIdx.x;
|
|
if (i >= batch_size) return;
|
|
|
|
int tid = threadIdx.x;
|
|
int label_i = labels[i];
|
|
|
|
|
|
extern __shared__ float s_query[];
|
|
for (int k = tid; k < dim; k += blockDim.x) {
|
|
s_query[k] = x[i * dim + k];
|
|
}
|
|
|
|
|
|
float* s_pos_hist = &s_query[dim];
|
|
float* s_neg_hist = &s_pos_hist[num_bins];
|
|
|
|
for (int k = tid; k < num_bins; k += blockDim.x) {
|
|
s_pos_hist[k] = 0.0f;
|
|
s_neg_hist[k] = 0.0f;
|
|
}
|
|
|
|
__syncthreads();
|
|
|
|
float width = max_dist / (float)num_bins;
|
|
float half_width = width * 0.5f;
|
|
|
|
for (int j = tid; j < batch_size; j += blockDim.x) {
|
|
if (i == j) continue;
|
|
|
|
float dist_sq = 0.0f;
|
|
const float* row_j = x + j * dim;
|
|
|
|
|
|
int k = 0;
|
|
for (; k + 3 < dim; k += 4) {
|
|
float d0 = s_query[k] - row_j[k];
|
|
float d1 = s_query[k+1] - row_j[k+1];
|
|
float d2 = s_query[k+2] - row_j[k+2];
|
|
float d3 = s_query[k+3] - row_j[k+3];
|
|
dist_sq += d0*d0 + d1*d1 + d2*d2 + d3*d3;
|
|
}
|
|
for (; k < dim; ++k) {
|
|
float d = s_query[k] - row_j[k];
|
|
dist_sq += d*d;
|
|
}
|
|
|
|
int label_j = labels[j];
|
|
bool is_pos = (label_i == label_j);
|
|
|
|
|
|
for (int m = 0; m < num_bins; ++m) {
|
|
float center = half_width + m * width;
|
|
float diff = fabsf(dist_sq - center);
|
|
float weight = 1.0f - diff / width;
|
|
if (weight > 0.0f) {
|
|
|
|
if (is_pos) {
|
|
atomicAdd(&s_pos_hist[m], weight);
|
|
} else {
|
|
atomicAdd(&s_neg_hist[m], weight);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
__syncthreads();
|
|
|
|
if (tid == 0) {
|
|
float pos_cdf = 0.0f;
|
|
float neg_cdf = 0.0f;
|
|
float ap = 0.0f;
|
|
|
|
|
|
float total_pos = 0.0f;
|
|
for (int m = 0; m < num_bins; ++m) {
|
|
total_pos += s_pos_hist[m];
|
|
}
|
|
|
|
|
|
total_pos += 1e-10f;
|
|
|
|
|
|
for (int m = 0; m < num_bins; ++m) {
|
|
float p_hist = s_pos_hist[m];
|
|
float n_hist = s_neg_hist[m];
|
|
|
|
pos_cdf += p_hist;
|
|
neg_cdf += n_hist;
|
|
|
|
float precision = pos_cdf / (pos_cdf + neg_cdf + 1e-10f);
|
|
float delta_recall = p_hist / total_pos;
|
|
|
|
ap += precision * delta_recall;
|
|
}
|
|
|
|
ap_out[i] = ap;
|
|
}
|
|
}
|
|
|
|
torch::Tensor launch_fast_ap(torch::Tensor x, torch::Tensor labels, int num_bins) {
|
|
auto batch_size = x.size(0);
|
|
auto dim = x.size(1);
|
|
auto ap_out = torch::empty({batch_size}, x.options());
|
|
|
|
int shared_mem = (dim + 2 * num_bins) * sizeof(float);
|
|
|
|
dim3 blocks(batch_size);
|
|
dim3 threads(256);
|
|
|
|
fast_ap_loss_kernel<<<blocks, threads, shared_mem>>>(
|
|
x.data_ptr<float>(),
|
|
labels.data_ptr<long>(),
|
|
ap_out.data_ptr<float>(),
|
|
batch_size,
|
|
dim,
|
|
num_bins,
|
|
4.0f
|
|
);
|
|
|
|
return ap_out;
|
|
}
|
|
"""
|
|
|
|
cpp_source = """
|
|
torch::Tensor launch_fast_ap(torch::Tensor x, torch::Tensor labels, int num_bins);
|
|
"""
|
|
|
|
fast_ap_module = load_inline(
|
|
name='fast_ap_loss_ext',
|
|
cpp_sources=cpp_source,
|
|
cuda_sources=cuda_source,
|
|
functions=['launch_fast_ap'],
|
|
verbose=False
|
|
)
|
|
|
|
|
|
class ModelNew(nn.Module):
|
|
def __init__(self, num_bins=10):
|
|
super(ModelNew, self).__init__()
|
|
self.num_bins = num_bins
|
|
self.op = fast_ap_module
|
|
|
|
def forward(self, x: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
|
|
x = F.normalize(x, p=2, dim=1)
|
|
|
|
ap_per_sample = self.op.launch_fast_ap(x.contiguous(), labels.contiguous(), self.num_bins)
|
|
|
|
return 1.0 - ap_per_sample.mean() |