forked from ccf-ai-infra/GPUCodeForces
Merge pull request 'finish ReflectionPad2d #23' (#47) from gsd123/GPUCodeForces:ReflectionPad2d into main
This commit is contained in:
commit
17a95d57ba
|
|
@ -0,0 +1,372 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
# -------------------------------------------------------------
|
||||
# 常量定义 (与你之前的代码一致)
|
||||
# -------------------------------------------------------------
|
||||
BATCH_SIZE = 32
|
||||
CHANNELS = 64
|
||||
HEIGHT = 32 # H_in
|
||||
WIDTH = 32 # W_in
|
||||
PADDING = (1, 1, 2, 0)
|
||||
BLOCK_DIM_X = 16
|
||||
BLOCK_DIM_Y = 16
|
||||
|
||||
|
||||
# -------------------------------------------------------------
|
||||
|
||||
class ModelNew(nn.Module):
|
||||
"""
|
||||
ReflectionPad2d 的高性能 CUDA 融合核函数实现
|
||||
(V3: 修复了 'contiguous' 运行时错误)
|
||||
"""
|
||||
|
||||
def __init__(self, padding):
|
||||
super().__init__()
|
||||
|
||||
if isinstance(padding, int):
|
||||
self.pad_L = padding
|
||||
self.pad_R = padding
|
||||
self.pad_T = padding
|
||||
self.pad_B = padding
|
||||
else:
|
||||
self.pad_L = padding[0]
|
||||
self.pad_R = padding[1]
|
||||
self.pad_T = padding[2]
|
||||
self.pad_B = padding[3]
|
||||
|
||||
self.block_dim_x = BLOCK_DIM_X
|
||||
self.block_dim_y = BLOCK_DIM_Y
|
||||
self._compile_cuda_kernel()
|
||||
|
||||
def _compile_cuda_kernel(self):
|
||||
|
||||
cpp_header = f"""
|
||||
#include <torch/extension.h>
|
||||
|
||||
// C++ 接口
|
||||
torch::Tensor reflection_pad2d_forward_cuda(
|
||||
torch::Tensor input,
|
||||
int pad_L, int pad_R,
|
||||
int pad_T, int pad_B
|
||||
);
|
||||
"""
|
||||
|
||||
cuda_source = f"""
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#define BLOCK_DIM_X {self.block_dim_x}
|
||||
#define BLOCK_DIM_Y {self.block_dim_y}
|
||||
|
||||
__device__ inline int reflect_idx(
|
||||
int j, int pad_before, int W_in
|
||||
) {{
|
||||
if (j < pad_before) {{
|
||||
return pad_before - j;
|
||||
}} else if (j < (pad_before + W_in)) {{
|
||||
return j - pad_before;
|
||||
}} else {{
|
||||
int j_rel = j - (pad_before + W_in);
|
||||
return W_in - 2 - j_rel;
|
||||
}}
|
||||
}}
|
||||
|
||||
__global__ void reflection_pad2d_fused_kernel(
|
||||
const float* __restrict__ input_data,
|
||||
float* __restrict__ output_data,
|
||||
int N, int C,
|
||||
int H_in, int W_in,
|
||||
int H_out, int W_out,
|
||||
int pad_L, int pad_R,
|
||||
int pad_T, int pad_B
|
||||
) {{
|
||||
extern __shared__ float s_in[];
|
||||
|
||||
const int n_idx = blockIdx.x;
|
||||
const int c_idx = blockIdx.y;
|
||||
const int tid_x = threadIdx.x;
|
||||
const int tid_y = threadIdx.y;
|
||||
|
||||
const float* p_in = input_data + (n_idx * C + c_idx) * (H_in * W_in);
|
||||
float* p_out = output_data + (n_idx * C + c_idx) * (H_out * W_out);
|
||||
|
||||
// Pass 1: Load to shared memory
|
||||
for (int i = tid_y; i < H_in; i += BLOCK_DIM_Y) {{
|
||||
for (int j = tid_x; j < W_in; j += BLOCK_DIM_X) {{
|
||||
s_in[i * W_in + j] = p_in[i * W_in + j];
|
||||
}}
|
||||
}}
|
||||
__syncthreads();
|
||||
|
||||
// Pass 2: Compute and store from shared memory
|
||||
for (int i = tid_y; i < H_out; i += BLOCK_DIM_Y) {{
|
||||
int in_i = reflect_idx(i, pad_T, H_in);
|
||||
|
||||
for (int j = tid_x; j < W_out; j += BLOCK_DIM_X) {{
|
||||
int in_j = reflect_idx(j, pad_L, W_in);
|
||||
p_out[i * W_out + j] = s_in[in_i * W_in + in_j];
|
||||
}}
|
||||
}}
|
||||
}}
|
||||
|
||||
// C++ 封装函数
|
||||
torch::Tensor reflection_pad2d_forward_cuda(
|
||||
torch::Tensor input,
|
||||
int pad_L, int pad_R,
|
||||
int pad_T, int pad_B
|
||||
) {{
|
||||
// 这个检查现在是安全的，因为我们在 Python 中确保了连续性
|
||||
TORCH_CHECK(input.is_contiguous(), "input must be contiguous");
|
||||
TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor");
|
||||
TORCH_CHECK(input.dim() == 4, "input must be 4D (N, C, H, W)");
|
||||
|
||||
const int64_t N_64 = input.size(0);
|
||||
const int64_t C_64 = input.size(1);
|
||||
const int64_t H_in_64 = input.size(2);
|
||||
const int64_t W_in_64 = input.size(3);
|
||||
|
||||
TORCH_CHECK(pad_L < W_in_64, "pad_L error");
|
||||
TORCH_CHECK(pad_R < W_in_64, "pad_R error");
|
||||
TORCH_CHECK(pad_T < H_in_64, "pad_T error");
|
||||
TORCH_CHECK(pad_B < H_in_64, "pad_B error");
|
||||
|
||||
const int64_t H_out_64 = H_in_64 + pad_T + pad_B;
|
||||
const int64_t W_out_64 = W_in_64 + pad_L + pad_R;
|
||||
|
||||
auto output = torch::empty({{N_64, C_64, H_out_64, W_out_64}}, input.options());
|
||||
|
||||
dim3 grid_dim(N_64, C_64);
|
||||
dim3 block_dim(BLOCK_DIM_X, BLOCK_DIM_Y);
|
||||
|
||||
const int shared_mem_size = H_in_64 * W_in_64 * sizeof(float);
|
||||
|
||||
reflection_pad2d_fused_kernel<<<grid_dim, block_dim, shared_mem_size>>>(
|
||||
input.data_ptr<float>(),
|
||||
output.data_ptr<float>(),
|
||||
static_cast<int>(N_64), static_cast<int>(C_64),
|
||||
static_cast<int>(H_in_64), static_cast<int>(W_in_64),
|
||||
static_cast<int>(H_out_64), static_cast<int>(W_out_64),
|
||||
pad_L, pad_R,
|
||||
pad_T, pad_B
|
||||
);
|
||||
|
||||
return output;
|
||||
}}
|
||||
"""
|
||||
|
||||
nvcc_flags = [
|
||||
'-O3',
|
||||
'--use_fast_math',
|
||||
'--expt-relaxed-constexpr'
|
||||
]
|
||||
|
||||
self.pad_op = load_inline(
|
||||
name="reflection_pad2d_op_v2_fixed", # (与 V2 编译的二进制文件相同)
|
||||
cpp_sources=cpp_header,
|
||||
cuda_sources=cuda_source,
|
||||
functions=["reflection_pad2d_forward_cuda"],
|
||||
extra_cuda_cflags=nvcc_flags,
|
||||
verbose=False
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
|
||||
# [修复]
|
||||
# 必须确保张量是连续的，才能传递给 C++/CUDA
|
||||
x_cont = x.contiguous()
|
||||
|
||||
# 调用我们编译好的 CUDA C++ 函数
|
||||
return self.pad_op.reflection_pad2d_forward_cuda(
|
||||
x_cont, # 传递连续的张量
|
||||
self.pad_L, self.pad_R,
|
||||
self.pad_T, self.pad_B
|
||||
)
|
||||
import torch
|
||||
|
||||
|
||||
import torch.nn as nn
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
# -------------------------------------------------------------
|
||||
# 常量定义 (与你之前的代码一致)
|
||||
# -------------------------------------------------------------
|
||||
BATCH_SIZE = 32
|
||||
CHANNELS = 64
|
||||
HEIGHT = 32 # H_in
|
||||
WIDTH = 32 # W_in
|
||||
PADDING = (1, 1, 2, 0)
|
||||
BLOCK_DIM_X = 16
|
||||
BLOCK_DIM_Y = 16
|
||||
|
||||
|
||||
# -------------------------------------------------------------
|
||||
|
||||
class ModelNew(nn.Module):
|
||||
"""
|
||||
ReflectionPad2d 的高性能 CUDA 融合核函数实现
|
||||
(V3: 修复了 'contiguous' 运行时错误)
|
||||
"""
|
||||
|
||||
def __init__(self, padding):
|
||||
super().__init__()
|
||||
|
||||
if isinstance(padding, int):
|
||||
self.pad_L = padding
|
||||
self.pad_R = padding
|
||||
self.pad_T = padding
|
||||
self.pad_B = padding
|
||||
else:
|
||||
self.pad_L = padding[0]
|
||||
self.pad_R = padding[1]
|
||||
self.pad_T = padding[2]
|
||||
self.pad_B = padding[3]
|
||||
|
||||
self.block_dim_x = BLOCK_DIM_X
|
||||
self.block_dim_y = BLOCK_DIM_Y
|
||||
self._compile_cuda_kernel()
|
||||
|
||||
def _compile_cuda_kernel(self):
|
||||
|
||||
cpp_header = f"""
|
||||
#include <torch/extension.h>
|
||||
|
||||
// C++ 接口
|
||||
torch::Tensor reflection_pad2d_forward_cuda(
|
||||
torch::Tensor input,
|
||||
int pad_L, int pad_R,
|
||||
int pad_T, int pad_B
|
||||
);
|
||||
"""
|
||||
|
||||
cuda_source = f"""
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#define BLOCK_DIM_X {self.block_dim_x}
|
||||
#define BLOCK_DIM_Y {self.block_dim_y}
|
||||
|
||||
__device__ inline int reflect_idx(
|
||||
int j, int pad_before, int W_in
|
||||
) {{
|
||||
if (j < pad_before) {{
|
||||
return pad_before - j;
|
||||
}} else if (j < (pad_before + W_in)) {{
|
||||
return j - pad_before;
|
||||
}} else {{
|
||||
int j_rel = j - (pad_before + W_in);
|
||||
return W_in - 2 - j_rel;
|
||||
}}
|
||||
}}
|
||||
|
||||
__global__ void reflection_pad2d_fused_kernel(
|
||||
const float* __restrict__ input_data,
|
||||
float* __restrict__ output_data,
|
||||
int N, int C,
|
||||
int H_in, int W_in,
|
||||
int H_out, int W_out,
|
||||
int pad_L, int pad_R,
|
||||
int pad_T, int pad_B
|
||||
) {{
|
||||
extern __shared__ float s_in[];
|
||||
|
||||
const int n_idx = blockIdx.x;
|
||||
const int c_idx = blockIdx.y;
|
||||
const int tid_x = threadIdx.x;
|
||||
const int tid_y = threadIdx.y;
|
||||
|
||||
const float* p_in = input_data + (n_idx * C + c_idx) * (H_in * W_in);
|
||||
float* p_out = output_data + (n_idx * C + c_idx) * (H_out * W_out);
|
||||
|
||||
// Pass 1: Load to shared memory
|
||||
for (int i = tid_y; i < H_in; i += BLOCK_DIM_Y) {{
|
||||
for (int j = tid_x; j < W_in; j += BLOCK_DIM_X) {{
|
||||
s_in[i * W_in + j] = p_in[i * W_in + j];
|
||||
}}
|
||||
}}
|
||||
__syncthreads();
|
||||
|
||||
// Pass 2: Compute and store from shared memory
|
||||
for (int i = tid_y; i < H_out; i += BLOCK_DIM_Y) {{
|
||||
int in_i = reflect_idx(i, pad_T, H_in);
|
||||
|
||||
for (int j = tid_x; j < W_out; j += BLOCK_DIM_X) {{
|
||||
int in_j = reflect_idx(j, pad_L, W_in);
|
||||
p_out[i * W_out + j] = s_in[in_i * W_in + in_j];
|
||||
}}
|
||||
}}
|
||||
}}
|
||||
|
||||
// C++ 封装函数
|
||||
torch::Tensor reflection_pad2d_forward_cuda(
|
||||
torch::Tensor input,
|
||||
int pad_L, int pad_R,
|
||||
int pad_T, int pad_B
|
||||
) {{
|
||||
// 这个检查现在是安全的，因为我们在 Python 中确保了连续性
|
||||
TORCH_CHECK(input.is_contiguous(), "input must be contiguous");
|
||||
TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor");
|
||||
TORCH_CHECK(input.dim() == 4, "input must be 4D (N, C, H, W)");
|
||||
|
||||
const int64_t N_64 = input.size(0);
|
||||
const int64_t C_64 = input.size(1);
|
||||
const int64_t H_in_64 = input.size(2);
|
||||
const int64_t W_in_64 = input.size(3);
|
||||
|
||||
TORCH_CHECK(pad_L < W_in_64, "pad_L error");
|
||||
TORCH_CHECK(pad_R < W_in_64, "pad_R error");
|
||||
TORCH_CHECK(pad_T < H_in_64, "pad_T error");
|
||||
TORCH_CHECK(pad_B < H_in_64, "pad_B error");
|
||||
|
||||
const int64_t H_out_64 = H_in_64 + pad_T + pad_B;
|
||||
const int64_t W_out_64 = W_in_64 + pad_L + pad_R;
|
||||
|
||||
auto output = torch::empty({{N_64, C_64, H_out_64, W_out_64}}, input.options());
|
||||
|
||||
dim3 grid_dim(N_64, C_64);
|
||||
dim3 block_dim(BLOCK_DIM_X, BLOCK_DIM_Y);
|
||||
|
||||
const int shared_mem_size = H_in_64 * W_in_64 * sizeof(float);
|
||||
|
||||
reflection_pad2d_fused_kernel<<<grid_dim, block_dim, shared_mem_size>>>(
|
||||
input.data_ptr<float>(),
|
||||
output.data_ptr<float>(),
|
||||
static_cast<int>(N_64), static_cast<int>(C_64),
|
||||
static_cast<int>(H_in_64), static_cast<int>(W_in_64),
|
||||
static_cast<int>(H_out_64), static_cast<int>(W_out_64),
|
||||
pad_L, pad_R,
|
||||
pad_T, pad_B
|
||||
);
|
||||
|
||||
return output;
|
||||
}}
|
||||
"""
|
||||
|
||||
nvcc_flags = [
|
||||
'-O3',
|
||||
'--use_fast_math',
|
||||
'--expt-relaxed-constexpr'
|
||||
]
|
||||
|
||||
self.pad_op = load_inline(
|
||||
name="reflection_pad2d_op_v2_fixed", # (与 V2 编译的二进制文件相同)
|
||||
cpp_sources=cpp_header,
|
||||
cuda_sources=cuda_source,
|
||||
functions=["reflection_pad2d_forward_cuda"],
|
||||
extra_cuda_cflags=nvcc_flags,
|
||||
verbose=False
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
|
||||
# [修复]
|
||||
# 必须确保张量是连续的，才能传递给 C++/CUDA
|
||||
x_cont = x.contiguous()
|
||||
|
||||
# 调用我们编译好的 CUDA C++ 函数
|
||||
return self.pad_op.reflection_pad2d_forward_cuda(
|
||||
x_cont, # 传递连续的张量
|
||||
self.pad_L, self.pad_R,
|
||||
self.pad_T, self.pad_B
|
||||
)
|
||||
|
|
@ -0,0 +1,50 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
# -------------------------------------------------------------
|
||||
# 常量定义
|
||||
# -------------------------------------------------------------
|
||||
BATCH_SIZE = 32
|
||||
CHANNELS = 64
|
||||
HEIGHT = 32 # H_in
|
||||
WIDTH = 32 # W_in
|
||||
|
||||
# (pad_L, pad_R, pad_T, pad_B)
|
||||
PADDING = (1, 1, 2, 0)
|
||||
|
||||
|
||||
# -------------------------------------------------------------
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
nn.ReflectionPad2d 的纯 PyTorch 基准实现
|
||||
(使用 F.pad)
|
||||
"""
|
||||
|
||||
def __init__(self, padding):
|
||||
super().__init__()
|
||||
|
||||
if isinstance(padding, int):
|
||||
# F.pad 需要 (left, right, top, bottom) 格式
|
||||
self.padding_tuple = (padding, padding, padding, padding)
|
||||
else:
|
||||
self.padding_tuple = padding
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
# F.pad 的 padding 格式是 (pad_dim_0_left, pad_dim_0_right, ...)
|
||||
# 对应 (N, C, H, W)，我们需要 pad 最后两个维度
|
||||
# F.pad 接受的顺序是 (pad_W_left, pad_W_right, pad_H_top, pad_H_bottom)
|
||||
return F.pad(x, self.padding_tuple, mode='reflect')
|
||||
|
||||
|
||||
def get_inputs():
|
||||
"""
|
||||
生成一个 (N, C, H, W) 形状的输入
|
||||
"""
|
||||
x = torch.randn(BATCH_SIZE, CHANNELS, HEIGHT, WIDTH, dtype=torch.float32)
|
||||
return [x]
|
||||
|
||||
|
||||
def get_init_inputs():
|
||||
return [PADDING]
|
||||
|
|
@ -0,0 +1,109 @@
|
|||
You write custom CUDA kernels to replace the PyTorch operators in the given EvoNorm 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 normalization+affine_transform+nonlinear_gating), or algorithmic changes (such as optimized memory access patterns). You are only limited by your imagination.
|
||||
|
||||
Here's a summary of the technologies used in the ReflectionPad2d CUDA implementation:
|
||||
|
||||
Key Technologies Used:
|
||||
|
||||
Inline CUDA Extension in PyTorch: Uses torch.utils.cpp_extension.load_inline() to compile and load CUDA code directly within Python, eliminating separate compilation steps.
|
||||
|
||||
Fused GPU Kernel Design: Implements a single kernel that combines data loading and padding operations into one efficient pass, minimizing kernel launch overhead.
|
||||
|
||||
Shared Memory Optimization: Leverages CUDA shared memory (s_in[]) to cache the entire input feature map for each channel, enabling fast data access compared to global memory.
|
||||
|
||||
Two-Phase Execution Strategy:
|
||||
|
||||
Pass 1: Loads input data from global memory to shared memory using grid-strided loops
|
||||
|
||||
Pass 2: Performs reflection padding calculations reading from shared memory and writing to global output
|
||||
|
||||
2D Thread Blocking: Employs 2D thread blocks (BLOCK_DIM_X/Y = 16) for efficient parallelization across height and width dimensions.
|
||||
|
||||
Mathematical Reflection Indexing: Implements a device-side reflect_idx function that calculates reflection indices using arithmetic operations rather than conditional branching for better performance.
|
||||
|
||||
Grid-Strided Loops: Uses strided loops in both loading and computation phases to handle arbitrary tensor sizes while maintaining load balancing.
|
||||
|
||||
Batched Channel Processing: Processes multiple batches and channels concurrently through 2D grid dimensions (grid_dim(N, C)).
|
||||
|
||||
Memory Contiguity Enforcement: Explicitly ensures input tensor contiguity in Python (x.contiguous()) before passing to CUDA, with runtime validation.
|
||||
|
||||
Comprehensive Error Checking: Includes extensive bounds checking for padding values and tensor dimensions to ensure valid operations.
|
||||
|
||||
Performance Optimizations:
|
||||
|
||||
Shared memory caching of entire input feature maps
|
||||
|
||||
Coalesced memory access patterns
|
||||
|
||||
Single synchronization point between loading and computation phases
|
||||
|
||||
Compiler optimizations (-O3, --use_fast_math)
|
||||
|
||||
Grid-strided loops for efficient workload distribution
|
||||
|
||||
Restricted pointers for better compiler optimization
|
||||
|
||||
Architecture Features:
|
||||
|
||||
Separate C++ header and CUDA source code organization
|
||||
|
||||
Template-style parameter passing for padding values
|
||||
|
||||
Automatic output tensor allocation with correct dimensions
|
||||
|
||||
Efficient memory stride calculations for 4D tensor layout (N, C, H, W)
|
||||
|
||||
|
||||
|
||||
|
||||
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
|
||||
|
||||
# -------------------------------------------------------------
|
||||
# 常量定义
|
||||
# -------------------------------------------------------------
|
||||
BATCH_SIZE = 32
|
||||
CHANNELS = 64
|
||||
HEIGHT = 32 # H_in
|
||||
WIDTH = 32 # W_in
|
||||
|
||||
# (pad_L, pad_R, pad_T, pad_B)
|
||||
PADDING = (1, 1, 2, 0)
|
||||
|
||||
|
||||
# -------------------------------------------------------------
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
nn.ReflectionPad2d 的纯 PyTorch 基准实现
|
||||
(使用 F.pad)
|
||||
"""
|
||||
|
||||
def __init__(self, padding):
|
||||
super().__init__()
|
||||
|
||||
if isinstance(padding, int):
|
||||
# F.pad 需要 (left, right, top, bottom) 格式
|
||||
self.padding_tuple = (padding, padding, padding, padding)
|
||||
else:
|
||||
self.padding_tuple = padding
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
# F.pad 的 padding 格式是 (pad_dim_0_left, pad_dim_0_right, ...)
|
||||
# 对应 (N, C, H, W)，我们需要 pad 最后两个维度
|
||||
# F.pad 接受的顺序是 (pad_W_left, pad_W_right, pad_H_top, pad_H_bottom)
|
||||
return F.pad(x, self.padding_tuple, mode='reflect')
|
||||
|
||||
|
||||
def get_inputs():
|
||||
"""
|
||||
生成一个 (N, C, H, W) 形状的输入
|
||||
"""
|
||||
x = torch.randn(BATCH_SIZE, CHANNELS, HEIGHT, WIDTH, dtype=torch.float32)
|
||||
return [x]
|
||||
|
||||
|
||||
def get_init_inputs():
|
||||
return [PADDING]
|
||||
|
|
@ -0,0 +1,77 @@
|
|||
###########################################################
|
||||
# 性能和精度验证程序
|
||||
###########################################################
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import time
|
||||
from ReflectionPad2d_torch import Model, get_inputs, get_init_inputs
|
||||
from ReflectionPad2d_cuda 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 torch.relu 平均执行时间: {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()
|
||||
Loading…
Reference in New Issue