1. 使用 PyTorch 实现 GPU 加速的卷积滤波(如边缘检测)
python
import torch
import torch.nn as nn
import cv2
import numpy as np
# 检查 GPU 是否可用
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
# 读取图像并转换为 PyTorch 张量
image = cv2.imread("input.jpg") # 读取 BGR 格式图像
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # 转为 RGB
image_tensor = torch.from_numpy(image).float().permute(2, 0, 1) # HWC -> CHW
image_tensor = image_tensor.unsqueeze(0).to(device) # 添加 batch 维度并移至 GPU
# 定义边缘检测卷积核(Sobel算子)
conv_layer = nn.Conv2d(
in_channels=3,
out_channels=3,
kernel_size=3,
bias=False,
padding=1
).to(device)
# 手动设置 Sobel 核权重(示例,仅作用于水平边缘)
sobel_kernel = torch.tensor([
[[[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]]], # Red 通道
[[[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]]], # Green 通道
[[[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]]], # Blue 通道
], dtype=torch.float32).repeat(3, 1, 1, 1).to(device)
conv_layer.weight.data = sobel_kernel
# 执行卷积操作(GPU加速)
with torch.no_grad():
output_tensor = conv_layer(image_tensor)
# 将结果转换回 numpy 并保存
output = output_tensor.squeeze(0).permute(1, 2, 0).cpu().numpy()
output = np.clip(output, 0, 255).astype(np.uint8)
cv2.imwrite("edge_detection_gpu.jpg", cv2.cvtColor(output, cv2.COLOR_RGB2BGR))
2. 使用 OpenCV 的 CUDA 模块加速高斯模糊
python
import cv2
import time
# 检查 OpenCV 是否支持 CUDA
print("CUDA devices:", cv2.cuda.getCudaEnabledDeviceCount())
# 读取图像并上传到 GPU
image = cv2.imread("input.jpg")
gpu_image = cv2.cuda_GpuMat()
gpu_image.upload(image)
# 创建 GPU 加速的高斯滤波器
gaussian_filter = cv2.cuda.createGaussianFilter(
cv2.CV_8UC3, # 输入类型 (8-bit unsigned, 3 channels)
cv2.CV_8UC3, # 输出类型
(15, 15), # 核大小
0 # Sigma(自动计算)
)
# 执行滤波(重复多次测试速度)
start_time = time.time()
for _ in range(100): # 重复 100 次模拟大数据量
gpu_blur = gaussian_filter.apply(gpu_image)
end_time = time.time()
# 下载结果到 CPU 并保存
result = gpu_blur.download()
print(f"GPU Time: {end_time - start_time:.4f} seconds")
cv2.imwrite("blur_gpu.jpg", result)
3. 使用 CuPy 加速图像傅里叶变换
python
import cupy as cp
import cv2
import numpy as np
import time
# 读取图像并转为灰度
image = cv2.imread("input.jpg", cv2.IMREAD_GRAYSCALE)
# 将 numpy 数组转为 CuPy 数组(上传到 GPU)
image_gpu = cp.asarray(image)
# 快速傅里叶变换(FFT)和逆变换(IFFT)
start_time = time.time()
fft_gpu = cp.fft.fft2(image_gpu)
fft_shift = cp.fft.fftshift(fft_gpu)
magnitude_spectrum = cp.log(cp.abs(fft_shift))
end_time = time.time()
# 将结果转回 CPU
magnitude_cpu = cp.asnumpy(magnitude_spectrum)
print(f"GPU FFT Time: {end_time - start_time:.4f} seconds")
# 归一化并保存频谱图
magnitude_cpu = cv2.normalize(magnitude_cpu, None, 0, 255, cv2.NORM_MINMAX)
cv2.imwrite("fft_spectrum_gpu.jpg", magnitude_cpu.astype(np.uint8))
4. 使用 Numba 编写自定义 GPU 核函数(图像反色)
python
from numba import cuda
import numpy as np
import cv2
import time
# 读取图像
image = cv2.imread("input.jpg")
height, width, channels = image.shape
# 定义 GPU 核函数
@cuda.jit
def invert_colors_kernel(image):
x, y = cuda.grid(2)
if x < image.shape[0] and y < image.shape[1]:
for c in range(3): # 遍历 RGB 通道
image[x, y, c] = 255 - image[x, y, c]
# 将图像上传到 GPU
image_gpu = cuda.to_device(image)
# 配置线程和块
threads_per_block = (16, 16)
blocks_per_grid_x = (height + threads_per_block[0] - 1) // threads_per_block[0]
blocks_per_grid_y = (width + threads_per_block[1] - 1) // threads_per_block[1]
blocks_per_grid = (blocks_per_grid_x, blocks_per_grid_y)
# 执行核函数
start_time = time.time()
invert_colors_kernel[blocks_per_grid, threads_per_block](image_gpu)
cuda.synchronize() # 等待 GPU 完成
end_time = time.time()
# 下载结果并保存
image_cpu = image_gpu.copy_to_host()
print(f"GPU Invert Time: {end_time - start_time:.6f} seconds")
cv2.imwrite("inverted_gpu.jpg", image_cpu)
5. 使用 PyTorch 实现实时风格迁移(GPU加速)
python
import torch
import torchvision.models as models
from torchvision import transforms
from PIL import Image
# 加载预训练模型到 GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = models.vgg19(pretrained=True).features.to(device).eval()
# 图像预处理
preprocess = transforms.Compose([
transforms.Resize(512),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
# 加载内容图像和风格图像
content_image = Image.open("content.jpg")
style_image = Image.open("style.jpg")
# 将图像转为张量并移至 GPU
content_tensor = preprocess(content_image).unsqueeze(0).to(device)
style_tensor = preprocess(style_image).unsqueeze(0).to(device)
# 定义风格迁移函数(示例,需完整实现损失计算和优化)
def style_transfer(model, content_input, style_input, iterations=500):
# 创建可优化图像
input_image = content_input.clone().requires_grad_(True)
# 定义优化器
optimizer = torch.optim.LBFGS([input_image])
# 风格迁移循环
for i in range(iterations):
def closure():
optimizer.zero_grad()
# 提取特征并计算损失(需实现具体细节)
# ...
return total_loss
optimizer.step(closure)
return input_image
# 执行风格迁移(需补充完整代码)
output_image = style_transfer(model, content_tensor, style_tensor)
# 后处理并保存结果
output_image = output_image.squeeze().cpu().detach()
output_image = transforms.ToPILImage()(output_image)
output_image.save("style_transfer_gpu.jpg")
关键说明
-
硬件依赖:需 NVIDIA GPU 并安装正确版本的 CUDA 和 cuDNN。
-
库安装 :
bashpip install torch torchvision opencv-python-headless cupy numba
-
性能对比:与 CPU 版本相比,GPU 加速通常快 10-100 倍(取决于任务复杂度)。
-
适用场景 :
- PyTorch:适合深度学习相关的图像处理(如 GAN、超分辨率)。
- OpenCV CUDA:适合传统图像处理加速(滤波、特征提取)。
- CuPy/Numba:适合自定义数值计算或科研算法。