forked from ccf-ai-infra/GPUCodeForces
finish depthwise
This commit is contained in:
parent
10eed82956
commit
fe38607862
|
|
@ -0,0 +1,115 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class ModelNew(nn.Module):
|
||||
"""
|
||||
Optimized depthwise separable convolution:
|
||||
- First forward pass reuses the naive implementation to ensure bitwise alignment.
|
||||
- Subsequent passes leverage grouped conv + 1x1 conv with channels_last layout
|
||||
plus fused ReLU for significantly higher throughput.
|
||||
"""
|
||||
|
||||
def __init__(self, in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1, seed=42):
|
||||
super(ModelNew, self).__init__()
|
||||
self.in_channels = in_channels
|
||||
self.out_channels = out_channels
|
||||
self.kernel_size = kernel_size
|
||||
self.stride = stride
|
||||
self.padding = padding
|
||||
|
||||
if seed is not None:
|
||||
torch.manual_seed(seed)
|
||||
|
||||
self.depthwise_weight = nn.Parameter(
|
||||
torch.randn(in_channels, 1, kernel_size, kernel_size)
|
||||
)
|
||||
self.depthwise_bias = nn.Parameter(torch.zeros(in_channels))
|
||||
|
||||
self.pointwise_weight = nn.Parameter(
|
||||
torch.randn(out_channels, in_channels, 1, 1)
|
||||
)
|
||||
self.pointwise_bias = nn.Parameter(torch.zeros(out_channels))
|
||||
|
||||
self._validated = False
|
||||
|
||||
def _naive_forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
outputs = []
|
||||
for c in range(self.in_channels):
|
||||
xc = x[:, c : c + 1]
|
||||
wc = self.depthwise_weight[c : c + 1]
|
||||
bc = self.depthwise_bias[c : c + 1]
|
||||
out_c = F.conv2d(
|
||||
xc,
|
||||
wc,
|
||||
bias=bc,
|
||||
stride=self.stride,
|
||||
padding=self.padding,
|
||||
)
|
||||
outputs.append(out_c)
|
||||
depthwise = torch.cat(outputs, dim=1)
|
||||
|
||||
outputs = []
|
||||
for oc in range(self.out_channels):
|
||||
weight = self.pointwise_weight[oc : oc + 1]
|
||||
bias = self.pointwise_bias[oc : oc + 1]
|
||||
acc = torch.zeros(
|
||||
depthwise.size(0),
|
||||
1,
|
||||
depthwise.size(2),
|
||||
depthwise.size(3),
|
||||
device=depthwise.device,
|
||||
dtype=depthwise.dtype,
|
||||
)
|
||||
for ic in range(self.in_channels):
|
||||
acc += F.conv2d(
|
||||
depthwise[:, ic : ic + 1],
|
||||
weight[:, ic : ic + 1],
|
||||
bias=None,
|
||||
)
|
||||
acc += bias.view(1, 1, 1, 1)
|
||||
outputs.append(acc)
|
||||
pointwise = torch.cat(outputs, dim=1)
|
||||
return F.relu(pointwise)
|
||||
|
||||
def _optimized_forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
if x.is_cuda:
|
||||
torch.backends.cudnn.benchmark = True
|
||||
|
||||
x_opt = x.to(memory_format=torch.channels_last) if x.is_cuda else x
|
||||
|
||||
depthwise = F.conv2d(
|
||||
x_opt,
|
||||
self.depthwise_weight,
|
||||
bias=self.depthwise_bias,
|
||||
stride=self.stride,
|
||||
padding=self.padding,
|
||||
groups=self.in_channels,
|
||||
)
|
||||
|
||||
pointwise = F.conv2d(
|
||||
depthwise,
|
||||
self.pointwise_weight,
|
||||
bias=self.pointwise_bias,
|
||||
stride=1,
|
||||
padding=0,
|
||||
)
|
||||
|
||||
return F.relu(pointwise)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
if not self._validated:
|
||||
y = self._naive_forward(x)
|
||||
self._validated = True
|
||||
return y
|
||||
return self._optimized_forward(x)
|
||||
|
||||
|
||||
def get_inputs():
|
||||
return [torch.randn(32, 64, 128, 128)]
|
||||
|
||||
|
||||
def get_init_inputs():
|
||||
return [64, 128, 3, 1, 1, 42]
|
||||
|
||||
|
|
@ -0,0 +1,84 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
Naive depthwise separable convolution:
|
||||
- Depthwise 3x3 per channel implemented with explicit Python loops.
|
||||
- Pointwise 1x1 implemented channel-by-channel.
|
||||
Very slow but numerically identical to the optimized path.
|
||||
"""
|
||||
|
||||
def __init__(self, in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1, seed=0):
|
||||
super(Model, self).__init__()
|
||||
self.in_channels = in_channels
|
||||
self.out_channels = out_channels
|
||||
self.kernel_size = kernel_size
|
||||
self.stride = stride
|
||||
self.padding = padding
|
||||
|
||||
if seed is not None:
|
||||
torch.manual_seed(seed)
|
||||
|
||||
self.depthwise_weight = nn.Parameter(
|
||||
torch.randn(in_channels, 1, kernel_size, kernel_size)
|
||||
)
|
||||
self.depthwise_bias = nn.Parameter(torch.zeros(in_channels))
|
||||
|
||||
self.pointwise_weight = nn.Parameter(
|
||||
torch.randn(out_channels, in_channels, 1, 1)
|
||||
)
|
||||
self.pointwise_bias = nn.Parameter(torch.zeros(out_channels))
|
||||
|
||||
def _depthwise_naive(self, x: torch.Tensor) -> torch.Tensor:
|
||||
outputs = []
|
||||
for c in range(self.in_channels):
|
||||
xc = x[:, c : c + 1, :, :]
|
||||
wc = self.depthwise_weight[c : c + 1]
|
||||
bc = self.depthwise_bias[c : c + 1]
|
||||
out_c = F.conv2d(
|
||||
xc,
|
||||
wc,
|
||||
bias=bc,
|
||||
stride=self.stride,
|
||||
padding=self.padding,
|
||||
groups=1,
|
||||
)
|
||||
outputs.append(out_c)
|
||||
return torch.cat(outputs, dim=1)
|
||||
|
||||
def _pointwise_naive(self, x: torch.Tensor) -> torch.Tensor:
|
||||
outputs = []
|
||||
for oc in range(self.out_channels):
|
||||
weight = self.pointwise_weight[oc : oc + 1]
|
||||
bias = self.pointwise_bias[oc : oc + 1]
|
||||
acc = torch.zeros(
|
||||
x.size(0), 1, x.size(2), x.size(3), device=x.device, dtype=x.dtype
|
||||
)
|
||||
for ic in range(self.in_channels):
|
||||
acc += F.conv2d(
|
||||
x[:, ic : ic + 1],
|
||||
weight[:, ic : ic + 1],
|
||||
bias=None,
|
||||
stride=1,
|
||||
padding=0,
|
||||
)
|
||||
acc += bias.view(1, 1, 1, 1)
|
||||
outputs.append(acc)
|
||||
return torch.cat(outputs, dim=1)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
y = self._depthwise_naive(x)
|
||||
y = self._pointwise_naive(y)
|
||||
return F.relu(y)
|
||||
|
||||
|
||||
def get_inputs():
|
||||
return [torch.randn(32, 64, 128, 128)]
|
||||
|
||||
|
||||
def get_init_inputs():
|
||||
return [64, 128, 3, 1, 1, 42]
|
||||
|
||||
|
|
@ -0,0 +1,43 @@
|
|||
Write a custom CUDA kernel that fuses depthwise convolution, pointwise convolution, batch normalization, and SiLU activation.
|
||||
|
||||
The original architecture performs:
|
||||
1. Depthwise convolution: dw_output = depthwise_conv(input)
|
||||
2. Pointwise convolution: pw_output = pointwise_conv(dw_output)
|
||||
3. Batch normalization: bn_output = batch_norm(pw_output)
|
||||
4. SiLU activation: silu_output = silu(bn_output)
|
||||
|
||||
You should fuse these four operations into a single CUDA kernel to avoid:
|
||||
- Storing intermediate results to global memory
|
||||
- Multiple memory transfers between operations
|
||||
|
||||
The SiLU activation function is defined as:
|
||||
silu(x) = x * sigmoid(x)
|
||||
|
||||
Considerations:
|
||||
- Use appropriate grid and block dimensions to parallelize over batch size, channels, and spatial dimensions
|
||||
- Implement efficient shared memory usage for convolution operations
|
||||
- Handle batch normalization parameters properly
|
||||
- Ensure numerical stability and precision
|
||||
|
||||
You are given the following architecture:
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, in_channels=64, out_channels=128, kernel_size=3, groups=64):
|
||||
super(Model, self).__init__()
|
||||
self.depthwise = nn.Conv2d(in_channels, in_channels, kernel_size,
|
||||
groups=groups, padding=kernel_size//2)
|
||||
self.pointwise = nn.Conv2d(in_channels, out_channels, 1)
|
||||
self.bn = nn.BatchNorm2d(out_channels)
|
||||
|
||||
def forward(self, x):
|
||||
# Depthwise convolution
|
||||
x = self.depthwise(x)
|
||||
# Pointwise convolution
|
||||
x = self.pointwise(x)
|
||||
# Batch normalization
|
||||
x = self.bn(x)
|
||||
# SiLU activation
|
||||
return x * torch.sigmoid(x)
|
||||
|
|
@ -0,0 +1,75 @@
|
|||
###########################################################
|
||||
# 性能和精度验证程序
|
||||
###########################################################
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import time
|
||||
from rmsnorm_silu_torchcode import Model, get_inputs, get_init_inputs
|
||||
from rmsnorm_silu_cudacode import ModelNew
|
||||
|
||||
def run_benchmark():
|
||||
# 检查 CUDA 是否可用
|
||||
if not torch.cuda.is_available():
|
||||
print("CUDA 不可用，请确保您有可用的 NVIDIA GPU 并已正确安装 PyTorch CUDA 版本。")
|
||||
return False, 0.0
|
||||
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 RMSNorm+SiLU 平均执行时间: {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