forked from ccf-ai-infra/GPUCodeForces
101 lines
2.9 KiB
Python
101 lines
2.9 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
from torch.utils.cpp_extension import load_inline
|
|
|
|
|
|
class ModelNew(nn.Module):
|
|
def __init__(self, epsilon=1.0):
|
|
super().__init__()
|
|
self.epsilon = epsilon
|
|
self._compile_cuda_kernel()
|
|
|
|
def _compile_cuda_kernel(self):
|
|
cpp_source = """
|
|
torch::Tensor poly_loss_cuda(torch::Tensor logits, torch::Tensor labels, float epsilon);
|
|
"""
|
|
|
|
cuda_source = """
|
|
#include <torch/extension.h>
|
|
#include <cuda_runtime.h>
|
|
#include <math.h>
|
|
|
|
__global__ void poly_loss_kernel(
|
|
const float* __restrict__ Z,
|
|
const long* __restrict__ Y,
|
|
float* __restrict__ L_out,
|
|
const int batch_size,
|
|
const int num_classes,
|
|
const float epsilon)
|
|
{
|
|
const int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
|
const int stride = blockDim.x * gridDim.x;
|
|
|
|
if (tid >= batch_size) return;
|
|
|
|
const float* logit_row = Z + tid * num_classes;
|
|
const int target_class = (int)Y[tid];
|
|
|
|
float row_max = -3.402823466e+38F;
|
|
float sum_exp = 0.0f;
|
|
|
|
|
|
for (int j = 0; j < num_classes; ++j) {
|
|
row_max = fmaxf(row_max, logit_row[j]);
|
|
}
|
|
|
|
|
|
for (int j = 0; j < num_classes; ++j) {
|
|
sum_exp += expf(logit_row[j] - row_max);
|
|
}
|
|
|
|
float log_sum_exp = row_max + logf(sum_exp);
|
|
|
|
|
|
float logit_t = logit_row[target_class];
|
|
float p_t = expf(logit_t - log_sum_exp);
|
|
|
|
|
|
float l_ce = log_sum_exp - logit_t;
|
|
|
|
|
|
L_out[tid] = l_ce + epsilon * (1.0f - p_t);
|
|
}
|
|
|
|
torch::Tensor poly_loss_cuda(torch::Tensor logits, torch::Tensor labels, float epsilon) {
|
|
auto Z_c = logits.contiguous();
|
|
auto Y_c = labels.contiguous();
|
|
|
|
const int batch_size = Z_c.size(0);
|
|
const int num_classes = Z_c.size(1);
|
|
|
|
auto output = torch::empty({batch_size}, Z_c.options());
|
|
|
|
const int threads = 256;
|
|
const int max_blocks = 65535;
|
|
const int blocks = std::min(batch_size, max_blocks);
|
|
|
|
poly_loss_kernel<<<blocks, threads>>>(
|
|
Z_c.data_ptr<float>(),
|
|
Y_c.data_ptr<long>(),
|
|
output.data_ptr<float>(),
|
|
batch_size,
|
|
num_classes,
|
|
epsilon
|
|
);
|
|
|
|
return output;
|
|
}
|
|
"""
|
|
|
|
self.op = load_inline(
|
|
name="poly_loss_v1",
|
|
cpp_sources=cpp_source,
|
|
cuda_sources=cuda_source,
|
|
functions=["poly_loss_cuda"],
|
|
extra_cuda_cflags=["-O3", "--use_fast_math"],
|
|
verbose=False
|
|
)
|
|
|
|
def forward(self, logits, labels):
|
|
loss_elementwise = self.op.poly_loss_cuda(logits, labels, self.epsilon)
|
|
return loss_elementwise.mean() |