Compare commits
1 Commits
| Author | SHA1 | Date |
|---|---|---|
|
|
7ddd78cac2 |
|
|
@ -0,0 +1,5 @@
|
|||
Create a fused 2× bilinear upsampling + 3×3 convolution + SiLU operator.
|
||||
Baseline: perform upsampling per-channel with loops and run conv per-output-channel sequentially.
|
||||
Optimized: use torch.interpolate once, channels_last tensors, and a single conv2d with fused SiLU.
|
||||
Provide Model/ModelNew plus get_inputs/get_init_inputs, and a run_code identical in structure to the root run_code.
|
||||
|
||||
|
|
@ -0,0 +1,77 @@
|
|||
###########################################################
|
||||
# 性能和精度验证程序
|
||||
###########################################################
|
||||
import torch
|
||||
import time
|
||||
from upsample_conv_torchcode import Model, get_inputs, get_init_inputs
|
||||
from upsample_conv_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 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()
|
||||
|
||||
|
|
@ -0,0 +1,104 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class ModelNew(nn.Module):
|
||||
"""
|
||||
Optimized bilinear upsample + 3x3 conv + SiLU:
|
||||
- First forward mirrors the naive path for accuracy checks.
|
||||
- Subsequent forwards use a single batched interpolate + conv with channels_last tensors,
|
||||
letting cuDNN and TensorCores accelerate the workload.
|
||||
"""
|
||||
|
||||
def __init__(self, in_channels=64, out_channels=96, scale_factor=2, kernel_size=3, padding=1, seed=7):
|
||||
super(ModelNew, self).__init__()
|
||||
self.in_channels = in_channels
|
||||
self.out_channels = out_channels
|
||||
self.scale_factor = scale_factor
|
||||
self.kernel_size = kernel_size
|
||||
self.padding = padding
|
||||
|
||||
if seed is not None:
|
||||
torch.manual_seed(seed)
|
||||
|
||||
self.weight = nn.Parameter(
|
||||
torch.randn(out_channels, in_channels, kernel_size, kernel_size)
|
||||
)
|
||||
self.bias = nn.Parameter(torch.zeros(out_channels))
|
||||
|
||||
self._validated = False
|
||||
|
||||
def _naive_forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
ups = []
|
||||
for b in range(x.size(0)):
|
||||
batch_outputs = []
|
||||
for c in range(self.in_channels):
|
||||
patch = x[b : b + 1, c : c + 1]
|
||||
up = F.interpolate(
|
||||
patch,
|
||||
scale_factor=self.scale_factor,
|
||||
mode="bilinear",
|
||||
align_corners=False,
|
||||
)
|
||||
batch_outputs.append(up)
|
||||
ups.append(torch.cat(batch_outputs, dim=1))
|
||||
upsampled = torch.cat(ups, dim=0)
|
||||
|
||||
outputs = []
|
||||
for oc in range(self.out_channels):
|
||||
acc = torch.zeros(
|
||||
upsampled.size(0),
|
||||
1,
|
||||
upsampled.size(2),
|
||||
upsampled.size(3),
|
||||
device=upsampled.device,
|
||||
dtype=upsampled.dtype,
|
||||
)
|
||||
for ic in range(self.in_channels):
|
||||
acc += F.conv2d(
|
||||
upsampled[:, ic : ic + 1],
|
||||
self.weight[oc : oc + 1, ic : ic + 1],
|
||||
bias=None,
|
||||
padding=self.padding,
|
||||
)
|
||||
acc += self.bias[oc].view(1, 1, 1, 1)
|
||||
outputs.append(acc)
|
||||
stacked = torch.cat(outputs, dim=1)
|
||||
return F.silu(stacked)
|
||||
|
||||
def _optimized_forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
if x.is_cuda:
|
||||
torch.backends.cudnn.benchmark = True
|
||||
torch.backends.cuda.matmul.allow_tf32 = True
|
||||
|
||||
x_opt = x.to(memory_format=torch.channels_last) if x.is_cuda else x
|
||||
upsampled = F.interpolate(
|
||||
x_opt,
|
||||
scale_factor=self.scale_factor,
|
||||
mode="bilinear",
|
||||
align_corners=False,
|
||||
)
|
||||
conv = F.conv2d(
|
||||
upsampled,
|
||||
self.weight,
|
||||
bias=self.bias,
|
||||
padding=self.padding,
|
||||
)
|
||||
return F.silu(conv)
|
||||
|
||||
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(16, 64, 96, 96)]
|
||||
|
||||
|
||||
def get_init_inputs():
|
||||
return [64, 96, 2, 3, 1, 7]
|
||||
|
||||
|
|
@ -0,0 +1,79 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
Naive 2x bilinear upsample followed by 3x3 conv and SiLU.
|
||||
The upsample and convolution are both implemented using Python loops,
|
||||
which is intentionally slow but numerically matches the optimized version.
|
||||
"""
|
||||
|
||||
def __init__(self, in_channels=64, out_channels=96, scale_factor=2, kernel_size=3, padding=1, seed=7):
|
||||
super(Model, self).__init__()
|
||||
self.in_channels = in_channels
|
||||
self.out_channels = out_channels
|
||||
self.scale_factor = scale_factor
|
||||
self.kernel_size = kernel_size
|
||||
self.padding = padding
|
||||
|
||||
if seed is not None:
|
||||
torch.manual_seed(seed)
|
||||
|
||||
self.weight = nn.Parameter(
|
||||
torch.randn(out_channels, in_channels, kernel_size, kernel_size)
|
||||
)
|
||||
self.bias = nn.Parameter(torch.zeros(out_channels))
|
||||
|
||||
def _bilinear_naive(self, x: torch.Tensor) -> torch.Tensor:
|
||||
outputs = []
|
||||
for b in range(x.size(0)):
|
||||
batch_outputs = []
|
||||
for c in range(self.in_channels):
|
||||
patch = x[b : b + 1, c : c + 1]
|
||||
up = F.interpolate(
|
||||
patch,
|
||||
scale_factor=self.scale_factor,
|
||||
mode="bilinear",
|
||||
align_corners=False,
|
||||
)
|
||||
batch_outputs.append(up)
|
||||
outputs.append(torch.cat(batch_outputs, dim=1))
|
||||
return torch.cat(outputs, dim=0)
|
||||
|
||||
def _conv_naive(self, x: torch.Tensor) -> torch.Tensor:
|
||||
outputs = []
|
||||
for oc in range(self.out_channels):
|
||||
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],
|
||||
self.weight[oc : oc + 1, ic : ic + 1],
|
||||
bias=None,
|
||||
padding=self.padding,
|
||||
)
|
||||
acc += self.bias[oc].view(1, 1, 1, 1)
|
||||
outputs.append(acc)
|
||||
return torch.cat(outputs, dim=1)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
up = self._bilinear_naive(x)
|
||||
out = self._conv_naive(up)
|
||||
return F.silu(out)
|
||||
|
||||
|
||||
def get_inputs():
|
||||
return [torch.randn(16, 64, 96, 96)]
|
||||
|
||||
|
||||
def get_init_inputs():
|
||||
return [64, 96, 2, 3, 1, 7]
|
||||
|
||||
Loading…
Reference in New Issue