forked from ccf-ai-infra/GPUCodeForces
96 lines
3.2 KiB
Python
96 lines
3.2 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
from torch.utils.cpp_extension import load_inline
|
|
|
|
|
|
class ModelNew(nn.Module):
|
|
def __init__(self, num_bins=10, min_val=0.0, max_val=1.0):
|
|
super().__init__()
|
|
self.num_bins = num_bins
|
|
self.min_val = min_val
|
|
self.max_val = max_val
|
|
self.step = (max_val - min_val) / num_bins
|
|
self._compile_cuda_kernel()
|
|
|
|
def _compile_cuda_kernel(self):
|
|
cpp_source = """
|
|
torch::Tensor histogramloss_cuda(torch::Tensor pos, torch::Tensor neg, int num_bins, float min_val, float max_val);
|
|
"""
|
|
|
|
cuda_source = """
|
|
#include <torch/extension.h>
|
|
#include <cuda_runtime.h>
|
|
|
|
__global__ void histogram_kernel(
|
|
const float* __restrict__ data,
|
|
float* __restrict__ hist,
|
|
const int n_elements,
|
|
const int num_bins,
|
|
const float min_val,
|
|
const float step)
|
|
{
|
|
const int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
|
const int stride = blockDim.x * gridDim.x;
|
|
|
|
for (int i = tid; i < n_elements; i += stride) {
|
|
float val = data[i];
|
|
for (int b = 0; b < num_bins; ++b) {
|
|
float center = min_val + (b + 0.5f) * step;
|
|
float diff = fabsf(val - center);
|
|
float weight = fmaxf(0.0f, 1.0f - diff / step);
|
|
atomicAdd(&hist[b], weight);
|
|
}
|
|
}
|
|
}
|
|
|
|
torch::Tensor histogramloss_cuda(torch::Tensor pos, torch::Tensor neg, int num_bins, float min_val, float max_val) {
|
|
auto pos_c = pos.contiguous();
|
|
auto neg_c = neg.contiguous();
|
|
|
|
auto pos_hist = torch::zeros({num_bins}, pos.options());
|
|
auto neg_hist = torch::zeros({num_bins}, neg.options());
|
|
|
|
float step = (max_val - min_val) / num_bins;
|
|
|
|
const int threads = 256;
|
|
|
|
int pos_blocks = min((int)((pos_c.numel() + threads - 1) / threads), 65535);
|
|
histogram_kernel<<<pos_blocks, threads>>>(
|
|
pos_c.data_ptr<float>(),
|
|
pos_hist.data_ptr<float>(),
|
|
pos_c.numel(),
|
|
num_bins,
|
|
min_val,
|
|
step
|
|
);
|
|
|
|
int neg_blocks = min((int)((neg_c.numel() + threads - 1) / threads), 65535);
|
|
histogram_kernel<<<neg_blocks, threads>>>(
|
|
neg_c.data_ptr<float>(),
|
|
neg_hist.data_ptr<float>(),
|
|
neg_c.numel(),
|
|
num_bins,
|
|
min_val,
|
|
step
|
|
);
|
|
|
|
auto pos_cdf = torch::cumsum(pos_hist, 0);
|
|
pos_cdf = pos_cdf / (pos_cdf[num_bins - 1] + 1e-8);
|
|
|
|
auto neg_pdf = neg_hist / (neg_hist.sum() + 1e-8);
|
|
|
|
return (neg_pdf * pos_cdf).sum();
|
|
}
|
|
"""
|
|
|
|
self.op = load_inline(
|
|
name="histogramloss_op",
|
|
cpp_sources=cpp_source,
|
|
cuda_sources=cuda_source,
|
|
functions=["histogramloss_cuda"],
|
|
extra_cuda_cflags=["-O3"],
|
|
verbose=False
|
|
)
|
|
|
|
def forward(self, pos, neg):
|
|
return self.op.histogramloss_cuda(pos, neg, self.num_bins, self.min_val, self.max_val) |