forked from ccf-ai-infra/GPUCodeForces
88 lines
2.1 KiB
Python
88 lines
2.1 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>
|
|
|
|
__global__ void f_divergence_kernel(
|
|
const float* __restrict__ p,
|
|
const float* __restrict__ q,
|
|
float* __restrict__ output,
|
|
int n,
|
|
float eps
|
|
) {
|
|
extern __shared__ float sdata[];
|
|
int tid = threadIdx.x;
|
|
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
|
|
|
float local_sum = 0.0f;
|
|
for (int i = idx; i < n; i += blockDim.x * gridDim.x) {
|
|
float p_val = p[i];
|
|
float q_val = q[i];
|
|
float diff = p_val - q_val;
|
|
local_sum += (diff * diff) / (q_val + eps);
|
|
}
|
|
|
|
sdata[tid] = local_sum;
|
|
__syncthreads();
|
|
|
|
for (int s = blockDim.x / 2; s > 0; s >>= 1) {
|
|
if (tid < s) {
|
|
sdata[tid] += sdata[tid + s];
|
|
}
|
|
__syncthreads();
|
|
}
|
|
|
|
if (tid == 0) {
|
|
atomicAdd(output, sdata[0]);
|
|
}
|
|
}
|
|
|
|
torch::Tensor f_divergence_cuda(torch::Tensor p, torch::Tensor q, float eps) {
|
|
int batch_size = p.size(0);
|
|
int num_classes = p.size(1);
|
|
int n = batch_size * num_classes;
|
|
|
|
auto output = at::zeros({1}, p.options());
|
|
|
|
int threads = 256;
|
|
int blocks = 128;
|
|
int shared_mem = threads * sizeof(float);
|
|
|
|
f_divergence_kernel<<<blocks, threads, shared_mem>>>(
|
|
p.data_ptr<float>(),
|
|
q.data_ptr<float>(),
|
|
output.data_ptr<float>(),
|
|
n,
|
|
eps
|
|
);
|
|
|
|
return output / batch_size;
|
|
}
|
|
"""
|
|
|
|
cpp_source = """
|
|
torch::Tensor f_divergence_cuda(torch::Tensor p, torch::Tensor q, float eps);
|
|
"""
|
|
|
|
f_divergence_loss = load_inline(
|
|
name="f_divergence_loss",
|
|
cpp_sources=cpp_source,
|
|
cuda_sources=cuda_source,
|
|
functions=["f_divergence_cuda"],
|
|
verbose=False
|
|
)
|
|
|
|
|
|
class ModelNew(nn.Module):
|
|
def __init__(self, eps=1e-8):
|
|
super(ModelNew, self).__init__()
|
|
self.eps = eps
|
|
|
|
def forward(self, p, q):
|
|
p_prob = F.softmax(p, dim=1)
|
|
q_prob = F.softmax(q, dim=1)
|
|
return f_divergence_loss.f_divergence_cuda(p_prob, q_prob, self.eps) |