forked from ccf-ai-infra/GPUCodeForces
Merge pull request 'feat add a selu_clip1 #4' (#196) from zizi05/GPUCodeForces:selu_clip1 into main
This commit is contained in:
commit
25ecf13fb2
|
|
@ -0,0 +1,113 @@
|
|||
You write custom CUDA kernels to replace the PyTorch operators in the given SELU-Clip activation 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 the combined SELU calculation and clamp operators with a custom CUDA kernel (considering operator fusion opportunities to combine the element-wise SELU computation and range clipping into a single kernel) or adjust algorithms for better performance. You are only limited by your imagination.
|
||||
|
||||
Here's an example to show you the syntax of inline embedding custom CUDA operators in PyTorch (for reference of code structure, not functional alignment):
|
||||
The example given architecture (sample structure):
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
class Model(nn.Module):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
def forward(self, a, b):
|
||||
return a + b
|
||||
def get_inputs():
|
||||
# randomly generate input tensors based on the model architecture
|
||||
a = torch.randn(1, 128).cuda()
|
||||
b = torch.randn(1, 128).cuda()
|
||||
return [a, b]
|
||||
def get_init_inputs():
|
||||
# randomly generate tensors required for initialization based on the model architecture
|
||||
return []
|
||||
|
||||
The example new arch with custom CUDA kernels (sample structure):
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
# Define custom CUDA kernel and load it inline
|
||||
custom_add_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
__global__ void custom_add_kernel(const float* a, const float* b, float* out, int size) {
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx < size) {
|
||||
out[idx] = a[idx] + b[idx];
|
||||
}
|
||||
}
|
||||
torch::Tensor custom_add_cuda(torch::Tensor a, torch::Tensor b) {
|
||||
auto size = a.numel();
|
||||
auto out = torch::empty_like(a);
|
||||
const int block_size = 256;
|
||||
int num_blocks = (size + block_size - 1) / block_size;
|
||||
custom_add_kernel<<<num_blocks, block_size>>>(a.data_ptr<float>(), b.data_ptr<float>(), out.data_ptr<float>(), size);
|
||||
return out;
|
||||
}
|
||||
"""
|
||||
custom_add_cpp_source = "torch::Tensor custom_add_cuda(torch::Tensor a, torch::Tensor b);"
|
||||
custom_add = load_inline(
|
||||
name="custom_add",
|
||||
cpp_sources=custom_add_cpp_source,
|
||||
cuda_sources=custom_add_source,
|
||||
functions=["custom_add_cuda"],
|
||||
verbose=True
|
||||
)
|
||||
class Model(nn.Module):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.custom_add = custom_add
|
||||
def forward(self, a, b):
|
||||
return self.custom_add.custom_add_cuda(a, b)
|
||||
def get_inputs():
|
||||
# randomly generate input tensors based on the model architecture
|
||||
a = torch.randn(1, 128).cuda()
|
||||
b = torch.randn(1, 128).cuda()
|
||||
return [a, b]
|
||||
def get_init_inputs():
|
||||
# randomly generate tensors required for initialization based on the model architecture
|
||||
return []
|
||||
|
||||
You are given the following SELU-Clip activation architecture (base PyTorch implementation):
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
SELU-Clip activation function: Mathematical formulation is clamp(scale * (x if x > 0 else alpha*(exp(x)-1)), clip_min, clip_max),
|
||||
where alpha=1.6732632423543772 (SELU shape parameter), scale=1.0507009873554804 (SELU scaling parameter),
|
||||
clip_min=-5.0 (minimum clipping value), clip_max=5.0 (maximum clipping value).
|
||||
"""
|
||||
def __init__(self, alpha=1.6732632423543772, scale=1.0507009873554804, clip_min=-5.0, clip_max=5.0):
|
||||
super(Model, self).__init__()
|
||||
self.alpha = alpha
|
||||
self.scale = scale
|
||||
self.clip_min = clip_min
|
||||
self.clip_max = clip_max
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Applies SELU-Clip activation to the input tensor.
|
||||
Args:
|
||||
x (torch.Tensor): Input tensor with fixed shape (batch_size, dim)
|
||||
where batch_size=1024 and dim=2048.
|
||||
Returns:
|
||||
torch.Tensor: Output tensor with SELU-Clip applied, same shape as input.
|
||||
"""
|
||||
# Calculate SELU
|
||||
selu_out = self.scale * torch.where(
|
||||
x > 0,
|
||||
x,
|
||||
self.alpha * (torch.exp(x) - 1)
|
||||
)
|
||||
# Apply clipping
|
||||
clipped_out = torch.clamp(selu_out, self.clip_min, self.clip_max)
|
||||
return clipped_out
|
||||
|
||||
batch_size = 1024
|
||||
dim = 2048
|
||||
def get_inputs():
|
||||
# Randomly generate input tensor matching the fixed shape (batch_size, dim)
|
||||
x = torch.randn(batch_size, dim)
|
||||
return [x]
|
||||
def get_init_inputs():
|
||||
# Provide initialization parameters for the model (non-trainable hyperparameters)
|
||||
return (1.6732632423543772, 1.0507009873554804, -5.0, 5.0)
|
||||
|
|
@ -0,0 +1,74 @@
|
|||
###########################################################
|
||||
# 性能和精度验证程序
|
||||
###########################################################
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import time
|
||||
from selu_clip_torchcode import Model,get_inputs,get_init_inputs
|
||||
from selu_clip_cudacode import ModelNew
|
||||
|
||||
def run_benchmark():
|
||||
# 检查 CUDA 是否可用
|
||||
if not torch.cuda.is_available():
|
||||
print("CUDA 不可用，请确保您有可用的 NVIDIA GPU 并已正确安装 PyTorch CUDA 版本。")
|
||||
return
|
||||
else:
|
||||
device = torch.device("cuda")
|
||||
|
||||
# 初始化模型
|
||||
init_inputs = get_init_inputs()
|
||||
init_inputs = [
|
||||
x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in init_inputs
|
||||
]
|
||||
inputs = get_inputs()
|
||||
inputs = [
|
||||
x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in inputs
|
||||
]
|
||||
|
||||
torch_model = Model(*init_inputs).cuda()
|
||||
cuda_model = ModelNew(*init_inputs).cuda()
|
||||
|
||||
torch_model.eval()
|
||||
cuda_model.eval()
|
||||
|
||||
print("-------------------- 精度对齐验证 --------------------")
|
||||
with torch.no_grad():
|
||||
output_torch = torch_model( *inputs)
|
||||
output_cuda = cuda_model(*inputs)
|
||||
|
||||
precision_flag = torch.allclose(output_torch, output_cuda,rtol=1e-03)
|
||||
if precision_flag:
|
||||
print("✅ 精度对齐：两个模型的输出结果非常接近。")
|
||||
else:
|
||||
print("❌ 精度不一致！")
|
||||
|
||||
print("\n-------------------- 性能加速比测试 --------------------")
|
||||
num_iterations = 100
|
||||
|
||||
# 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 selu_clip 平均执行时间: {torch_time:.6f} 秒")
|
||||
print(f"自定义 CUDA 内核 平均执行时间: {cuda_time:.6f} 秒")
|
||||
speedup = 0
|
||||
if cuda_time > 0:
|
||||
speedup = torch_time / cuda_time
|
||||
print(f"加速比 (Speedup): {speedup:.2f}x")
|
||||
else:
|
||||
print("CUDA 内核执行时间为0，无法计算加速比。")
|
||||
return precision_flag,speedup
|
||||
if __name__ == "__main__":
|
||||
precision_flag,speedup = run_benchmark()
|
||||
|
|
@ -0,0 +1,153 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.autograd import Function
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
# 定义CUDA内核和C++绑定代码
|
||||
cuda_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <vector>
|
||||
|
||||
// SELU-Clip前向传播CUDA内核
|
||||
__global__ void selu_clip_forward_kernel(
|
||||
const float* __restrict__ x,
|
||||
float* __restrict__ output,
|
||||
float alpha,
|
||||
float scale,
|
||||
float clip_min,
|
||||
float clip_max,
|
||||
int num_elements) {
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx < num_elements) {
|
||||
float val = x[idx];
|
||||
// 计算SELU
|
||||
float selu_val;
|
||||
if (val > 0.0f) {
|
||||
selu_val = scale * val;
|
||||
} else {
|
||||
selu_val = scale * alpha * (expf(val) - 1.0f);
|
||||
}
|
||||
// 裁剪操作
|
||||
if (selu_val < clip_min) {
|
||||
output[idx] = clip_min;
|
||||
} else if (selu_val > clip_max) {
|
||||
output[idx] = clip_max;
|
||||
} else {
|
||||
output[idx] = selu_val;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// SELU-Clip反向传播CUDA内核
|
||||
__global__ void selu_clip_backward_kernel(
|
||||
const float* __restrict__ x,
|
||||
const float* __restrict__ grad_output,
|
||||
float* __restrict__ grad_input,
|
||||
float alpha,
|
||||
float scale,
|
||||
int num_elements) {
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx < num_elements) {
|
||||
float val = x[idx];
|
||||
float grad = grad_output[idx];
|
||||
// 计算梯度
|
||||
if (val > 0.0f) {
|
||||
grad_input[idx] = scale * grad;
|
||||
} else {
|
||||
grad_input[idx] = scale * alpha * expf(val) * grad;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// C++绑定前向函数
|
||||
torch::Tensor selu_clip_forward(
|
||||
torch::Tensor x,
|
||||
float alpha,
|
||||
float scale,
|
||||
float clip_min,
|
||||
float clip_max) {
|
||||
auto output = torch::empty_like(x);
|
||||
int num_elements = x.numel();
|
||||
int block_size = 256;
|
||||
int grid_size = (num_elements + block_size - 1) / block_size;
|
||||
|
||||
selu_clip_forward_kernel<<<grid_size, block_size>>>(
|
||||
x.data_ptr<float>(),
|
||||
output.data_ptr<float>(),
|
||||
alpha,
|
||||
scale,
|
||||
clip_min,
|
||||
clip_max,
|
||||
num_elements
|
||||
);
|
||||
return output;
|
||||
}
|
||||
|
||||
// C++绑定反向函数
|
||||
torch::Tensor selu_clip_backward(
|
||||
torch::Tensor x,
|
||||
torch::Tensor grad_output,
|
||||
float alpha,
|
||||
float scale) {
|
||||
auto grad_input = torch::empty_like(x);
|
||||
int num_elements = x.numel();
|
||||
int block_size = 256;
|
||||
int grid_size = (num_elements + block_size - 1) / block_size;
|
||||
|
||||
selu_clip_backward_kernel<<<grid_size, block_size>>>(
|
||||
x.data_ptr<float>(),
|
||||
grad_output.data_ptr<float>(),
|
||||
grad_input.data_ptr<float>(),
|
||||
alpha,
|
||||
scale,
|
||||
num_elements
|
||||
);
|
||||
return grad_input;
|
||||
}
|
||||
|
||||
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
||||
m.def("forward", &selu_clip_forward, "SELU-Clip forward");
|
||||
m.def("backward", &selu_clip_backward, "SELU-Clip backward");
|
||||
}
|
||||
"""
|
||||
|
||||
# 动态编译CUDA内核
|
||||
selu_clip_cuda = load_inline(
|
||||
name="selu_clip_cuda",
|
||||
cpp_sources=[],
|
||||
cuda_sources=[cuda_source],
|
||||
extra_cuda_cflags=["-O2"],
|
||||
with_cuda=True,
|
||||
verbose=False
|
||||
)
|
||||
|
||||
class SELUClipFunction(Function):
|
||||
"""自定义CUDA实现的SELU-Clip算子函数"""
|
||||
@staticmethod
|
||||
def forward(ctx, x, alpha, scale, clip_min, clip_max):
|
||||
ctx.alpha = alpha
|
||||
ctx.scale = scale
|
||||
ctx.save_for_backward(x)
|
||||
return selu_clip_cuda.forward(x, alpha, scale, clip_min, clip_max)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_output):
|
||||
x, = ctx.saved_tensors
|
||||
alpha = ctx.alpha
|
||||
scale = ctx.scale
|
||||
grad_input = selu_clip_cuda.backward(x, grad_output, alpha, scale)
|
||||
return grad_input, None, None, None, None
|
||||
|
||||
class ModelNew(nn.Module):
|
||||
"""使用内置CUDA内核的SELU-Clip模型"""
|
||||
def __init__(self, alpha=1.6732632423543772, scale=1.0507009873554804, clip_min=-5.0, clip_max=5.0):
|
||||
super(ModelNew, self).__init__()
|
||||
self.alpha = alpha
|
||||
self.scale = scale
|
||||
self.clip_min = clip_min
|
||||
self.clip_max = clip_max
|
||||
|
||||
def forward(self, x):
|
||||
return SELUClipFunction.apply(x, self.alpha, self.scale, self.clip_min, self.clip_max)
|
||||
|
|
@ -0,0 +1,32 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
class Model(nn.Module):
|
||||
"""PyTorch实现的SELU-Clip算子"""
|
||||
def __init__(self, alpha=1.6732632423543772, scale=1.0507009873554804, clip_min=-5.0, clip_max=5.0):
|
||||
super(Model, self).__init__()
|
||||
self.alpha = alpha
|
||||
self.scale = scale
|
||||
self.clip_min = clip_min
|
||||
self.clip_max = clip_max
|
||||
|
||||
def forward(self, x):
|
||||
# 计算SELU
|
||||
selu_out = self.scale * torch.where(
|
||||
x > 0,
|
||||
x,
|
||||
self.alpha * (torch.exp(x) - 1)
|
||||
)
|
||||
# 裁剪操作
|
||||
clipped_out = torch.clamp(selu_out, self.clip_min, self.clip_max)
|
||||
return clipped_out
|
||||
|
||||
def get_init_inputs():
|
||||
"""提供模型初始化参数"""
|
||||
return (1.6732632423543772, 1.0507009873554804, -5.0, 5.0)
|
||||
|
||||
def get_inputs():
|
||||
"""提供测试输入数据"""
|
||||
torch.manual_seed(42)
|
||||
x = torch.randn(1024, 2048) # 典型特征张量形状
|
||||
return (x,)
|
||||
Loading…
Reference in New Issue