fiexs localresponsenorm #8

S1/8/localresponsenorm_cuda.py

@@ -0,0 +1,98 @@
+# localresponsenorm_cuda.py
+import torch
+import torch.nn as nn
+from torch.utils.cpp_extension import load_inline
+
+from localresponsenorm_torch import N, C, H, W, SIZE, ALPHA, BETA, K
+
+class ModelNew(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self._compile_cuda_kernel()
+
+    def _compile_cuda_kernel(self):
+        half_size = SIZE // 2
+        
+        cpp_source = """
+        #include <torch/extension.h>
+        torch::Tensor lrn_forward_cuda(torch::Tensor input, int size, int half_size, float alpha, float beta, float k);
+        """
+
+        cuda_source = """
+        #include <cuda_runtime.h>
+        #include <device_launch_parameters.h>
+
+        #define BLOCK_SIZE 256
+
+        __global__ void lrn_fused_kernel(
+            const float* __restrict__ x,
+            float* __restrict__ y,
+            int N, int C, int H, int W,
+            int size, int half_size, float alpha, float beta, float k
+        ) {
+            int idx = blockIdx.x * blockDim.x + threadIdx.x;
+            if (idx >= N * C * H * W) return;
+
+            int n = idx / (C * H * W);
+            int rem = idx % (C * H * W);
+            int c = rem / (H * W);
+            int hw = rem % (H * W);
+
+            int base_idx = n * C * H * W + hw;
+            float sum_sq = 0.0f;
+            
+            int start_c = max(0, c - half_size);
+            int end_c = min(C - 1, c + half_size);
+
+            #pragma unroll 4
+            for (int window_c = start_c; window_c <= end_c; ++window_c) {
+                float val = x[base_idx + window_c * H * W];
+                sum_sq = __fmaf_rn(val, val, sum_sq);
+            }
+
+            float alpha_over_n = alpha / (float)size;
+            float scale = __fmaf_rn(alpha_over_n, sum_sq, k);
+            
+
+            float norm_factor = __powf(scale, beta);
+            float input_val = x[idx];
+            y[idx] = __fdividef(input_val, norm_factor);
+        }
+
+        torch::Tensor lrn_forward_cuda(torch::Tensor input, int size, int half_size, float alpha, float beta, float k) {
+            const int N = input.size(0);
+            const int C = input.size(1);
+            const int H = input.size(2);
+            const int W = input.size(3);
+            auto output = torch::empty_like(input);
+
+            const int n_elements = N * C * H * W;
+            const int blocks = (n_elements + BLOCK_SIZE - 1) / BLOCK_SIZE;
+
+            lrn_fused_kernel<<<blocks, BLOCK_SIZE>>>(
+                input.data_ptr<float>(),
+                output.data_ptr<float>(),
+                N, C, H, W,
+                size, half_size, alpha, beta, k
+            );
+            
+            return output;
+        }
+        """
+        
+        self.lrn_op = load_inline(
+            name="lrn_correct_formula",
+            cpp_sources=cpp_source, 
+            cuda_sources=cuda_source,
+            functions=["lrn_forward_cuda"],
+            extra_cuda_cflags=[
+                "-O3",
+                "--use_fast_math",
+                "--fmad=true"
+            ],
+            verbose=True
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        half_size = SIZE // 2
+        return self.lrn_op.lrn_forward_cuda(x.contiguous(), SIZE, half_size, ALPHA, BETA, K)

S1/8/localresponsenorm_torch.py

@@ -0,0 +1,27 @@
+# localresponsenorm_torch.py
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+N, C, H, W = 16, 64, 32, 32
+
+SIZE = 5      # 窗口大小 (n)
+ALPHA = 1e-4  # 缩放因子 (alpha)
+BETA = 0.75   # 幂指数 (beta)
+K = 2.0       # 偏置项 (k)
+
+class Model(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.lrn = nn.LocalResponseNorm(SIZE, alpha=ALPHA, beta=BETA, k=K)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.lrn(x)
+
+def get_inputs():
+    x = torch.randn(N, C, H, W, dtype=torch.float32)
+    return [x]
+
+def get_init_inputs():
+    return []

S1/8/run_code.py

@@ -0,0 +1,88 @@
+###########################################################
+# 性能和精度验证程序
+###########################################################
+import torch
+import torch.nn as nn
+import time
+from localresponsenorm_torch import Model, get_inputs, get_init_inputs
+from localresponsenorm_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)
+
+    # 更严格的精度检查
+    abs_diff = (output_torch - output_cuda).abs()
+    max_diff = abs_diff.max().item()
+    mean_diff = abs_diff.mean().item()
+    
+    print(f"最大差异: {max_diff:.6f}")
+    print(f"平均差异: {mean_diff:.6f}")
+    
+    precision_flag = torch.allclose(output_torch, output_cuda, rtol=1e-05, atol=1e-05)
+    if precision_flag:
+        print("✅ 精度对齐：两个模型的输出结果非常接近。")
+    else:
+        print("❌ 精度不一致！")
+        
+    print("\n-------------------- 性能加速比测试 --------------------")
+    num_iterations = 1000  # 增加迭代次数以获得更准确的时间测量
+    
+    # Warm up
+    for _ in range(100):
+        _ = 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 (matmul + relu) 平均执行时间: {torch_time:.6f} 秒")
+    print(f"自定义 CUDA ReLU 平均执行时间: {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()