forked from ccf-ai-infra/GPUCodeForces
Merge pull request 'finish FocalLoss #18' (#160) from gsd123/GPUCodeForces:gsd18 into main
This commit is contained in:
commit
5256224c5d
|
|
@ -0,0 +1,478 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
N, C, H, W = 32, 1, 64, 64
|
||||
|
||||
|
||||
class FocalLossCUDAOp(torch.autograd.Function):
|
||||
epsilon = 1e-8
|
||||
|
||||
def forward(ctx, input, target, alpha, gamma, reduction_id, op):
|
||||
if not input.is_cuda: input = input.cuda()
|
||||
if not target.is_cuda: target = target.cuda()
|
||||
|
||||
input = input.contiguous()
|
||||
target = target.contiguous()
|
||||
|
||||
n = input.numel()
|
||||
|
||||
output, partial_sums = op.focal_loss_forward_cuda(
|
||||
input,
|
||||
target,
|
||||
reduction_id,
|
||||
n,
|
||||
alpha,
|
||||
gamma,
|
||||
FocalLossCUDAOp.epsilon
|
||||
)
|
||||
|
||||
ctx.save_for_backward(input, target)
|
||||
ctx.reduction_id = reduction_id
|
||||
ctx.alpha = alpha
|
||||
ctx.gamma = gamma
|
||||
ctx.N = n
|
||||
ctx.op = op
|
||||
|
||||
return output
|
||||
|
||||
def backward(ctx, grad_output):
|
||||
input, target = ctx.saved_tensors
|
||||
|
||||
grad_out_scalar = 0.0
|
||||
grad_output_n = None
|
||||
|
||||
if ctx.reduction_id != 0:
|
||||
grad_out_scalar = grad_output[0]
|
||||
if ctx.reduction_id == 1:
|
||||
grad_out_scalar = grad_out_scalar / ctx.N
|
||||
else:
|
||||
grad_output_n = grad_output.contiguous()
|
||||
|
||||
grad_input = torch.empty_like(input)
|
||||
|
||||
ctx.op.focal_loss_backward_cuda(
|
||||
grad_out_scalar,
|
||||
grad_output_n,
|
||||
input,
|
||||
target,
|
||||
grad_input,
|
||||
ctx.reduction_id,
|
||||
ctx.N,
|
||||
ctx.alpha,
|
||||
ctx.gamma,
|
||||
FocalLossCUDAOp.epsilon
|
||||
)
|
||||
|
||||
return grad_input, None, None, None, None, None
|
||||
|
||||
|
||||
class ModelNew(nn.Module):
|
||||
def __init__(self, reduction='mean', alpha=0.25, gamma=2.0):
|
||||
super().__init__()
|
||||
self.alpha = float(alpha)
|
||||
self.gamma = float(gamma)
|
||||
|
||||
self.red_map = {'none': 0, 'mean': 1, 'sum': 2}
|
||||
if reduction not in self.red_map:
|
||||
raise ValueError("Invalid reduction")
|
||||
self.reduction_id = self.red_map[reduction]
|
||||
|
||||
self._compile_cuda_kernel()
|
||||
|
||||
def _compile_cuda_kernel(self):
|
||||
cpp_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <vector>
|
||||
|
||||
std::vector<torch::Tensor> focal_loss_forward_cuda(
|
||||
torch::Tensor input,
|
||||
torch::Tensor target,
|
||||
int reduction,
|
||||
int64_t n,
|
||||
float alpha,
|
||||
float gamma,
|
||||
float epsilon);
|
||||
|
||||
void focal_loss_backward_cuda(
|
||||
float grad_out_scalar,
|
||||
torch::Tensor grad_output_n,
|
||||
torch::Tensor input,
|
||||
torch::Tensor target,
|
||||
torch::Tensor grad_input,
|
||||
int reduction,
|
||||
int64_t n,
|
||||
float alpha,
|
||||
float gamma,
|
||||
float epsilon);
|
||||
"""
|
||||
|
||||
cuda_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <vector>
|
||||
#include <algorithm>
|
||||
#include <cuda_runtime.h>
|
||||
#include <cmath>
|
||||
#include <device_functions.h>
|
||||
#include <math_functions.h>
|
||||
|
||||
#define BLOCK_SIZE 256
|
||||
#define MAX_GRID_SIZE 4096
|
||||
|
||||
__inline__ __device__ double warp_reduce_sum_double(double val) {
|
||||
#pragma unroll
|
||||
for (int offset = 16; offset > 0; offset /= 2) {
|
||||
val += __shfl_down_sync(0xffffffff, val, offset);
|
||||
}
|
||||
return val;
|
||||
}
|
||||
|
||||
__inline__ __device__ double block_reduce_sum_double(double val) {
|
||||
__shared__ double shared[32];
|
||||
int lane = threadIdx.x % 32;
|
||||
int wid = threadIdx.x / 32;
|
||||
|
||||
val = warp_reduce_sum_double(val);
|
||||
if (lane == 0) shared[wid] = val;
|
||||
__syncthreads();
|
||||
|
||||
val = (threadIdx.x < blockDim.x / 32) ? shared[lane] : 0.0;
|
||||
if (wid == 0) val = warp_reduce_sum_double(val);
|
||||
return val;
|
||||
}
|
||||
|
||||
__global__ void reduce_kernel(
|
||||
const float* __restrict__ input,
|
||||
float* __restrict__ output,
|
||||
int n)
|
||||
{
|
||||
double local_sum = 0.0;
|
||||
int idx = threadIdx.x;
|
||||
|
||||
for (int i = idx; i < n; i += blockDim.x) {
|
||||
local_sum += (double)input[i];
|
||||
}
|
||||
|
||||
local_sum = block_reduce_sum_double(local_sum);
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
output[0] = (float)local_sum;
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void focal_fwd_kernel(
|
||||
const float* __restrict__ input,
|
||||
const float* __restrict__ target,
|
||||
float* __restrict__ output,
|
||||
int64_t n,
|
||||
int reduction,
|
||||
float* __restrict__ partial_sums,
|
||||
const double alpha_d,
|
||||
const double gamma_d,
|
||||
const double eps_d)
|
||||
{
|
||||
const int64_t n_vec = n / 4;
|
||||
const int64_t rem_start = n_vec * 4;
|
||||
|
||||
const int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int stride = blockDim.x * gridDim.x;
|
||||
|
||||
const float4* in_ptr = (const float4*)input;
|
||||
const float4* tgt_ptr = (const float4*)target;
|
||||
float4* out_ptr = (float4*)output;
|
||||
|
||||
double local_sum = 0.0;
|
||||
|
||||
for (int idx = i; idx < n_vec; idx += stride) {
|
||||
const float4 x_vec = in_ptr[idx];
|
||||
const float4 t_vec = tgt_ptr[idx];
|
||||
float4 loss_vec;
|
||||
|
||||
double loss[4];
|
||||
double x_d[4], t_d[4];
|
||||
|
||||
x_d[0] = (double)x_vec.x; t_d[0] = (double)t_vec.x;
|
||||
x_d[1] = (double)x_vec.y; t_d[1] = (double)t_vec.y;
|
||||
x_d[2] = (double)x_vec.z; t_d[2] = (double)t_vec.z;
|
||||
x_d[3] = (double)x_vec.w; t_d[3] = (double)t_vec.w;
|
||||
|
||||
#pragma unroll
|
||||
for(int k=0; k<4; ++k) {
|
||||
const double p_d = 1.0 / (1.0 + exp(-x_d[k]));
|
||||
const double p_t = p_d * t_d[k] + (1.0 - p_d) * (1.0 - t_d[k]);
|
||||
const double mod_factor = pow(1.0 - p_t, gamma_d);
|
||||
const double alpha_t = alpha_d * t_d[k] + (1.0 - alpha_d) * (1.0 - t_d[k]);
|
||||
|
||||
const double max_val = (x_d[k] > 0.0) ? x_d[k] : 0.0;
|
||||
const double stable_bce = max_val - x_d[k] * t_d[k] + log(1.0 + exp(-fabs(x_d[k])));
|
||||
|
||||
loss[k] = alpha_t * mod_factor * stable_bce;
|
||||
}
|
||||
|
||||
if (reduction == 0) {
|
||||
out_ptr[idx] = make_float4(
|
||||
(float)loss[0], (float)loss[1],
|
||||
(float)loss[2], (float)loss[3]
|
||||
);
|
||||
} else {
|
||||
local_sum += loss[0] + loss[1] + loss[2] + loss[3];
|
||||
}
|
||||
}
|
||||
|
||||
for (int idx = rem_start + i; idx < n; idx += stride) {
|
||||
const double x_d = (double)input[idx];
|
||||
const double t_d = (double)target[idx];
|
||||
|
||||
const double p_d = 1.0 / (1.0 + exp(-x_d));
|
||||
const double p_t = p_d * t_d + (1.0 - p_d) * (1.0 - t_d);
|
||||
const double mod_factor = pow(1.0 - p_t, gamma_d);
|
||||
const double alpha_t = alpha_d * t_d + (1.0 - alpha_d) * (1.0 - t_d);
|
||||
|
||||
const double max_val = (x_d > 0.0) ? x_d : 0.0;
|
||||
const double stable_bce = max_val - x_d * t_d + log(1.0 + exp(-fabs(x_d)));
|
||||
|
||||
const double loss = alpha_t * mod_factor * stable_bce;
|
||||
|
||||
if (reduction == 0) {
|
||||
output[idx] = (float)loss;
|
||||
} else {
|
||||
local_sum += loss;
|
||||
}
|
||||
}
|
||||
|
||||
if (reduction != 0) {
|
||||
local_sum = block_reduce_sum_double(local_sum);
|
||||
if (threadIdx.x == 0) {
|
||||
partial_sums[blockIdx.x] = (float)local_sum;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void focal_bwd_kernel(
|
||||
const double grad_out_scalar_d,
|
||||
const float* __restrict__ grad_output_n,
|
||||
const float* __restrict__ input,
|
||||
const float* __restrict__ target,
|
||||
float* __restrict__ grad_input,
|
||||
const int reduction,
|
||||
const int64_t n,
|
||||
const double alpha_d,
|
||||
const double gamma_d,
|
||||
const double eps_d)
|
||||
{
|
||||
const int64_t n_vec = n / 4;
|
||||
const int64_t rem_start = n_vec * 4;
|
||||
|
||||
const int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int stride = blockDim.x * gridDim.x;
|
||||
|
||||
const float4* in_ptr = (const float4*)input;
|
||||
const float4* tgt_ptr = (const float4*)target;
|
||||
const float4* grad_n_ptr = (const float4*)grad_output_n;
|
||||
float4* grad_in_ptr = (float4*)grad_input;
|
||||
|
||||
for (int idx = i; idx < n_vec; idx += stride) {
|
||||
const float4 x_vec = in_ptr[idx];
|
||||
const float4 t_vec = tgt_ptr[idx];
|
||||
|
||||
double grad_u[4];
|
||||
double x_d[4], t_d[4];
|
||||
|
||||
x_d[0] = (double)x_vec.x; t_d[0] = (double)t_vec.x;
|
||||
x_d[1] = (double)x_vec.y; t_d[1] = (double)t_vec.y;
|
||||
x_d[2] = (double)x_vec.z; t_d[2] = (double)t_vec.z;
|
||||
x_d[3] = (double)x_vec.w; t_d[3] = (double)t_vec.w;
|
||||
|
||||
double grad_out_d[4];
|
||||
if (reduction == 0) {
|
||||
const float4 grad_n_vec = grad_n_ptr[idx];
|
||||
grad_out_d[0] = (double)grad_n_vec.x;
|
||||
grad_out_d[1] = (double)grad_n_vec.y;
|
||||
grad_out_d[2] = (double)grad_n_vec.z;
|
||||
grad_out_d[3] = (double)grad_n_vec.w;
|
||||
} else {
|
||||
grad_out_d[0] = grad_out_scalar_d;
|
||||
grad_out_d[1] = grad_out_scalar_d;
|
||||
grad_out_d[2] = grad_out_scalar_d;
|
||||
grad_out_d[3] = grad_out_scalar_d;
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for(int k=0; k<4; ++k) {
|
||||
const double p_d = 1.0 / (1.0 + exp(-x_d[k]));
|
||||
const double dp_dx_d = p_d * (1.0 - p_d);
|
||||
|
||||
const double p_t_d = p_d * t_d[k] + (1.0 - p_d) * (1.0 - t_d[k]);
|
||||
const double one_minus_p_t = 1.0 - p_t_d;
|
||||
|
||||
const double alpha_t_d = alpha_d * t_d[k] + (1.0 - alpha_d) * (1.0 - t_d[k]);
|
||||
|
||||
const double mod_factor_d = pow(one_minus_p_t, gamma_d);
|
||||
const double mod_factor_g_minus_1_d = pow(one_minus_p_t, gamma_d - 1.0);
|
||||
|
||||
const double max_val = (x_d[k] > 0.0) ? x_d[k] : 0.0;
|
||||
const double stable_bce_d = max_val - x_d[k] * t_d[k] + log(1.0 + exp(-fabs(x_d[k])));
|
||||
|
||||
const double term1 = alpha_t_d * mod_factor_d * (p_d - t_d[k]);
|
||||
|
||||
const double d_p_t = (2.0 * t_d[k] - 1.0) * dp_dx_d;
|
||||
const double d_mod_factor = alpha_t_d * gamma_d * mod_factor_g_minus_1_d * (-d_p_t);
|
||||
|
||||
const double term2 = stable_bce_d * d_mod_factor;
|
||||
|
||||
grad_u[k] = (term1 + term2) * grad_out_d[k];
|
||||
}
|
||||
|
||||
grad_in_ptr[idx] = make_float4(
|
||||
(float)grad_u[0], (float)grad_u[1],
|
||||
(float)grad_u[2], (float)grad_u[3]
|
||||
);
|
||||
}
|
||||
|
||||
for (int idx = rem_start + i; idx < n; idx += stride) {
|
||||
const double x_d = (double)input[idx];
|
||||
const double t_d = (double)target[idx];
|
||||
|
||||
const double grad_out_d = (reduction == 0) ?
|
||||
(double)grad_output_n[idx] :
|
||||
grad_out_scalar_d;
|
||||
|
||||
const double p_d = 1.0 / (1.0 + exp(-x_d));
|
||||
const double dp_dx_d = p_d * (1.0 - p_d);
|
||||
const double p_t_d = p_d * t_d + (1.0 - p_d) * (1.0 - t_d);
|
||||
const double one_minus_p_t = 1.0 - p_t_d;
|
||||
const double alpha_t_d = alpha_d * t_d + (1.0 - alpha_d) * (1.0 - t_d);
|
||||
const double mod_factor_d = pow(one_minus_p_t, gamma_d);
|
||||
const double mod_factor_g_minus_1_d = pow(one_minus_p_t, gamma_d - 1.0);
|
||||
const double max_val = (x_d > 0.0) ? x_d : 0.0;
|
||||
const double stable_bce_d = max_val - x_d * t_d + log(1.0 + exp(-fabs(x_d)));
|
||||
const double term1 = alpha_t_d * mod_factor_d * (p_d - t_d);
|
||||
const double d_p_t = (2.0 * t_d - 1.0) * dp_dx_d;
|
||||
const double d_mod_factor = alpha_t_d * gamma_d * mod_factor_g_minus_1_d * (-d_p_t);
|
||||
const double term2 = stable_bce_d * d_mod_factor;
|
||||
|
||||
grad_input[idx] = (float)((term1 + term2) * grad_out_d);
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<torch::Tensor> focal_loss_forward_cuda(
|
||||
torch::Tensor input,
|
||||
torch::Tensor target,
|
||||
int reduction,
|
||||
int64_t n,
|
||||
float alpha,
|
||||
float gamma,
|
||||
float epsilon)
|
||||
{
|
||||
auto options = input.options();
|
||||
torch::Tensor output;
|
||||
torch::Tensor partial_sums;
|
||||
|
||||
const int block_size = BLOCK_SIZE;
|
||||
|
||||
// --- FIX IS HERE ---
|
||||
// Replaced std::min with ternary operator
|
||||
const int A_fwd = (int)((n / 4 + block_size - 1) / block_size);
|
||||
const int B_fwd = MAX_GRID_SIZE;
|
||||
const int grid_size = (A_fwd < B_fwd ? A_fwd : B_fwd);
|
||||
// --- END FIX ---
|
||||
|
||||
if (reduction == 0) {
|
||||
output = torch::empty_like(input);
|
||||
partial_sums = torch::empty({1}, options);
|
||||
} else {
|
||||
output = torch::zeros({1}, options);
|
||||
partial_sums = torch::zeros({grid_size}, options);
|
||||
}
|
||||
|
||||
focal_fwd_kernel<<<grid_size, block_size>>>(
|
||||
input.data_ptr<float>(),
|
||||
target.data_ptr<float>(),
|
||||
output.data_ptr<float>(),
|
||||
n,
|
||||
reduction,
|
||||
partial_sums.data_ptr<float>(),
|
||||
(double)alpha,
|
||||
(double)gamma,
|
||||
(double)epsilon
|
||||
);
|
||||
|
||||
if (reduction != 0) {
|
||||
reduce_kernel<<<1, BLOCK_SIZE>>>(
|
||||
partial_sums.data_ptr<float>(),
|
||||
output.data_ptr<float>(),
|
||||
grid_size
|
||||
);
|
||||
if (reduction == 1) {
|
||||
output.div_(n);
|
||||
}
|
||||
}
|
||||
|
||||
return {output, partial_sums};
|
||||
}
|
||||
|
||||
void focal_loss_backward_cuda(
|
||||
float grad_out_scalar,
|
||||
torch::Tensor grad_output_n,
|
||||
torch::Tensor input,
|
||||
torch::Tensor target,
|
||||
torch::Tensor grad_input,
|
||||
int reduction,
|
||||
int64_t n,
|
||||
float alpha,
|
||||
float gamma,
|
||||
float epsilon)
|
||||
{
|
||||
const int block_size = BLOCK_SIZE;
|
||||
|
||||
// --- FIX IS HERE ---
|
||||
// Replaced std::min with ternary operator
|
||||
const int A_bwd = (int)((n / 4 + block_size - 1) / block_size);
|
||||
const int B_bwd = MAX_GRID_SIZE;
|
||||
const int grid_size = (A_bwd < B_bwd ? A_bwd : B_bwd);
|
||||
// --- END FIX ---
|
||||
|
||||
const float* grad_output_n_ptr = (reduction == 0) ?
|
||||
grad_output_n.data_ptr<float>() :
|
||||
nullptr;
|
||||
|
||||
focal_bwd_kernel<<<grid_size, block_size>>>(
|
||||
(double)grad_out_scalar,
|
||||
grad_output_n_ptr,
|
||||
input.data_ptr<float>(),
|
||||
target.data_ptr<float>(),
|
||||
grad_input.data_ptr<float>(),
|
||||
reduction,
|
||||
n,
|
||||
(double)alpha,
|
||||
(double)gamma,
|
||||
(double)epsilon
|
||||
);
|
||||
}
|
||||
"""
|
||||
|
||||
self.op = load_inline(
|
||||
name='focal_loss_cuda_v2_vectorized_fix',
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=cuda_source,
|
||||
functions=['focal_loss_forward_cuda', 'focal_loss_backward_cuda'],
|
||||
extra_cuda_cflags=['-O3'],
|
||||
verbose=False
|
||||
)
|
||||
|
||||
def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
||||
if isinstance(input, (list, tuple)) and len(input) > 0:
|
||||
input = input[0]
|
||||
target = target[0] if len(target) > 0 else target
|
||||
|
||||
if target.dtype != input.dtype:
|
||||
target = target.to(input.dtype)
|
||||
|
||||
return FocalLossCUDAOp.apply(
|
||||
input,
|
||||
target,
|
||||
self.alpha,
|
||||
self.gamma,
|
||||
self.reduction_id,
|
||||
self.op
|
||||
)
|
||||
|
|
@ -0,0 +1,58 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
N, C, H, W = 32, 1, 64, 64
|
||||
|
||||
|
||||
class FocalLoss(nn.Module):
|
||||
def __init__(self, reduction='mean', alpha=0.25, gamma=2.0):
|
||||
super().__init__()
|
||||
self.reduction = reduction
|
||||
self.alpha = float(alpha)
|
||||
self.gamma = float(gamma)
|
||||
self.epsilon = 1e-8
|
||||
|
||||
def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
||||
|
||||
bce_loss = F.binary_cross_entropy_with_logits(input, target, reduction='none')
|
||||
|
||||
p = torch.sigmoid(input)
|
||||
|
||||
p_t = p * target + (1 - p) * (1 - target)
|
||||
|
||||
modulating_factor = (1.0 - p_t).pow(self.gamma)
|
||||
|
||||
alpha_t = self.alpha * target + (1 - self.alpha) * (1 - target)
|
||||
|
||||
loss = alpha_t * modulating_factor * bce_loss
|
||||
|
||||
if self.reduction == 'mean':
|
||||
return loss.mean()
|
||||
elif self.reduction == 'sum':
|
||||
return loss.sum()
|
||||
else:
|
||||
return loss
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, reduction='mean', alpha=0.25, gamma=2.0):
|
||||
super().__init__()
|
||||
self.op = FocalLoss(reduction, alpha, gamma)
|
||||
|
||||
def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
||||
if isinstance(input, (list, tuple)) and len(input) > 0:
|
||||
input = input[0]
|
||||
target = target[0] if len(target) > 0 else target
|
||||
|
||||
return self.op(input, target)
|
||||
|
||||
|
||||
def get_inputs():
|
||||
input = torch.randn(N, C, H, W, dtype=torch.float32)
|
||||
target = torch.randint(0, 2, (N, C, H, W), dtype=torch.float32)
|
||||
return [input, target]
|
||||
|
||||
|
||||
def get_init_inputs():
|
||||
return ['mean', 0.25, 2.0]
|
||||
|
|
@ -0,0 +1,111 @@
|
|||
You write custom CUDA kernels to replace the pytorch operators in the given GeGLU architecture to get speedups.
|
||||
|
||||
You have complete freedom to choose the set of operators you want to replace. You may make the decision to replace some operators with custom CUDA kernels and leave others unchanged. You may replace multiple operators with custom implementations, consider operator fusion opportunities (combining multiple operators into a single kernel, for example, combining chunk+gelu+elementwise_mul), or algorithmic changes (such as optimized memory access patterns). You are only limited by your imagination.
|
||||
|
||||
Core Optimization Techniques:
|
||||
|
||||
Performance Optimizations
|
||||
|
||||
Vectorized Memory Access - Uses float4 to process 4 elements simultaneously
|
||||
|
||||
Two-Stage Reduction - Warp-level + block-level shared memory reduction
|
||||
|
||||
Grid Size Optimization - Replaces std::min with ternary operator
|
||||
|
||||
Loop Unrolling - #pragma unroll for instruction-level parallelism
|
||||
|
||||
Memory Optimizations
|
||||
|
||||
Memory Coalescing - Ensures contiguous memory access patterns
|
||||
|
||||
Partial Sums Reduction - Distributed inter-block reduction strategy
|
||||
|
||||
Vectorized I/O - Uses float4 for both forward and backward passes
|
||||
|
||||
Numerical Stability
|
||||
|
||||
Stable BCE Computation - Log-sum-exp trick to prevent numerical overflow
|
||||
|
||||
Double Precision - Uses double for intermediate calculations
|
||||
|
||||
Epsilon Protection - Prevents division by zero errors
|
||||
|
||||
Focal Loss Specific Optimizations
|
||||
|
||||
Modulating Factor - Efficient computation of (1 - p_t)^gamma
|
||||
|
||||
Alpha Balancing - Class-balanced weighting calculation
|
||||
|
||||
Complete Gradient - Full gradient including both BCE and modulating terms
|
||||
|
||||
Kernel Design
|
||||
|
||||
Separate Paths - Distinct handling of reduction vs non-reduction paths
|
||||
|
||||
Remainder Processing - Efficient handling of elements after vectorization
|
||||
|
||||
Multi-Reduction Support - Full support for 'none', 'mean', and 'sum' reductions
|
||||
|
||||
These optimizations enable maximum performance for focal loss computation while maintaining numerical stability.
|
||||
|
||||
|
||||
|
||||
Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is:
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
N, C, H, W = 32, 1, 64, 64
|
||||
|
||||
|
||||
class FocalLoss(nn.Module):
|
||||
def __init__(self, reduction='mean', alpha=0.25, gamma=2.0):
|
||||
super().__init__()
|
||||
self.reduction = reduction
|
||||
self.alpha = float(alpha)
|
||||
self.gamma = float(gamma)
|
||||
self.epsilon = 1e-8
|
||||
|
||||
def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
||||
|
||||
bce_loss = F.binary_cross_entropy_with_logits(input, target, reduction='none')
|
||||
|
||||
p = torch.sigmoid(input)
|
||||
|
||||
p_t = p * target + (1 - p) * (1 - target)
|
||||
|
||||
modulating_factor = (1.0 - p_t).pow(self.gamma)
|
||||
|
||||
alpha_t = self.alpha * target + (1 - self.alpha) * (1 - target)
|
||||
|
||||
loss = alpha_t * modulating_factor * bce_loss
|
||||
|
||||
if self.reduction == 'mean':
|
||||
return loss.mean()
|
||||
elif self.reduction == 'sum':
|
||||
return loss.sum()
|
||||
else:
|
||||
return loss
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, reduction='mean', alpha=0.25, gamma=2.0):
|
||||
super().__init__()
|
||||
self.op = FocalLoss(reduction, alpha, gamma)
|
||||
|
||||
def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
||||
if isinstance(input, (list, tuple)) and len(input) > 0:
|
||||
input = input[0]
|
||||
target = target[0] if len(target) > 0 else target
|
||||
|
||||
return self.op(input, target)
|
||||
|
||||
|
||||
def get_inputs():
|
||||
input = torch.randn(N, C, H, W, dtype=torch.float32)
|
||||
target = torch.randint(0, 2, (N, C, H, W), dtype=torch.float32)
|
||||
return [input, target]
|
||||
|
||||
|
||||
def get_init_inputs():
|
||||
return ['mean', 0.25, 2.0]
|
||||
|
|
@ -0,0 +1,78 @@
|
|||
import torch
|
||||
import time
|
||||
from FocalLoss_torch import Model, get_inputs, get_init_inputs
|
||||
from FocalLoss_cuda import ModelNew
|
||||
|
||||
def run_benchmark():
|
||||
if not torch.cuda.is_available():
|
||||
print("CUDA 不可用")
|
||||
return
|
||||
|
||||
device = torch.device("cuda")
|
||||
|
||||
# 准备输入数据
|
||||
inputs = [x.cuda(device=device) for x in get_inputs()]
|
||||
init_inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_init_inputs()]
|
||||
|
||||
# 初始化模型
|
||||
torch_model = Model(*init_inputs).cuda()
|
||||
cuda_model = ModelNew(*init_inputs).cuda()
|
||||
|
||||
torch_model.eval()
|
||||
cuda_model.eval()
|
||||
|
||||
print("-------------------- 精度对齐验证 --------------------")
|
||||
with torch.no_grad():
|
||||
# 预热GPU
|
||||
_ = torch_model(*inputs)
|
||||
_ = cuda_model(*inputs)
|
||||
|
||||
# 正式测试
|
||||
output_torch = torch_model(*inputs)
|
||||
output_cuda = cuda_model(*inputs)
|
||||
|
||||
# 精度验证
|
||||
abs_diff = torch.abs(output_torch - output_cuda)
|
||||
max_diff = torch.max(abs_diff).item()
|
||||
mean_diff = torch.mean(abs_diff).item()
|
||||
|
||||
if max_diff < 1e-4 and mean_diff < 1e-5:
|
||||
print(f"✅ 精度对齐：最大误差 {max_diff:.6f}，平均误差 {mean_diff:.6f}")
|
||||
precision_flag = True
|
||||
else:
|
||||
print(f"❌ 精度不一致：最大误差 {max_diff:.6f}，平均误差 {mean_diff:.6f}")
|
||||
precision_flag = False
|
||||
|
||||
print("\n-------------------- 性能加速比测试 --------------------")
|
||||
num_iterations = 100
|
||||
|
||||
# 预热GPU
|
||||
for _ in range(10):
|
||||
_ = 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内置Swish平均执行时间: {torch_time:.6f}秒")
|
||||
print(f"自定义CUDA Swish平均执行时间: {cuda_time:.6f}秒")
|
||||
speedup = torch_time / cuda_time if cuda_time > 0 else 0
|
||||
print(f"加速比 (Speedup): {speedup:.2f}x")
|
||||
|
||||
return precision_flag, speedup
|
||||
|
||||
if __name__ == "__main__":
|
||||
precision_flag, speedup = run_benchmark()
|
||||
Loading…
Reference in New Issue