CUDA C++ 实现图像处理

最近课程里面有用到 NVIDIA CUDA 框架进行并行编程,实现了一些非常基本的图像处理的操作。
使用 CUDA 实现的并行加速能够极大的提升图像处理的效率,这也是为什么近几年的深度学习框架都要依托于 CUDA 进行计算加速。CUDA 本质上是 C/C++ 的拓展,因此对 C/C++ 熟悉的话上手也会很快。

读取保存图像

C++ 中读取图像和保存图像不像 Python 那样方便简单。一开始我想用 OpenCV 来读取/写入图像,但对于我写一个小的示例程序来说可谓麻烦了些,涉及到编译的问题。
后来看到了 lodepng 这个小巧方便的库,支持 C 和 C++,使用时只需要引用一个头文件就可以。读取和保存文件只需要如下的几行代码:

1
2
3
4
5
6
7
8
9
10
11
#include "lodepng.h"

// Read image
std::vector<unsigned char> image;
unsigned int width, height;

unsigned error = lodepng::decode(image, width, height, input_file);

// Save image
std::vector<unsigned char> out_image(image.size(), 255);
error = lodepng::encode(output_file, out_image, width, height);

加载好的图片像素有 4 个通道,按照 RGBARGBA... 的顺序在一维数组中排列的,其中 A 表示的是透明度,0 表示全透明,255 则完全显示当前像素。

CUDA Kernel

CUDA 并行部分的代码需要单独用 __global__ 关键词进行修饰,这一部分的代码会在 device 上执行。NVIDIA 自家的博文 An Even Easier Introduction to CUDA 就通俗易懂的介绍了 CUDA 的一些基本概念。CUDA C++ Programming Guide 有更为全面和详细的例子与参考。

下面是我实现的 RGB 转 Grayscale 的 Kernel 并行代码:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
// CUDA Kernel
__global__ void RGB2GrayKernel(unsigned char* input_image,
unsigned char* output_image,
int width, int height) {

// input_image size: width*height*Channels
// output_image size: width*height

int row = blockIdx.y * blockDim.y + threadIdx.y;
int col = blockIdx.x * blockDim.x + threadIdx.x;

if ((row < height) && (col < width)) {
// Offset in Grayscale image
int offset = row*width + col;

// Get RGB values
unsigned char r = input_image[offset * CHANNELS + 0];
unsigned char g = input_image[offset * CHANNELS + 1];
unsigned char b = input_image[offset * CHANNELS + 2];

// Convert to grayscale
output_image[offset] = 0.2126f*r + 0.7152f*g + 0.0722f*b;
}
}

通过 CUDA Kernel 全局变量 blockDimblockIdx 以及 threadIdx,可以得到当前运行 Kernel 的线程的编号,这样能够方便的访问部分数据。

在 CPU(host) 部分的 main 函数中,我们需要将图片数据拷贝到 device 上的 global memory 中,需要用到 cudaMalloc 以及 cudaMemcpy

1
2
3
4
5
6
7
8
// Allocate memory for CUDA device
size_t mem_size = width * height * sizeof(unsigned char);
unsigned char* dev_input, *dev_output;

cudaMalloc(&dev_input, mem_size * CHANNELS);
cudaMemcpy(dev_input, image.data(), mem_size * CHANNELS, cudaMemcpyHostToDevice);

cudaMalloc(&dev_output, mem_size);

main 中调用 CUDA Kernel 的时候,需要用到尖括号的表达。下面在一个 Block 里我用了 16x16 个线程,然后对于 width x height 的图像,可以得到所需的 Grid 的 Size,使得整个图片能被所有的线程都处理到。

1
2
3
4
5
6
7
#define BLOCK_SIZE 16

// Invoke CUDA kernel
dim3 dimBlock(BLOCK_SIZE, BLOCK_SIZE);
dim3 dimGrid(width / dimBlock.x, height / dimBlock.y);

RGB2GrayKernel<<<dimGrid, dimBlock>>>(dev_input, dev_output, width, height);

在调用完 Kernel 之后,我们需要将 CUDA device 内存中的数据拷贝到 CPU 上 main 程序里,这样可以在后续使用 lodepng 将图片保存。最后不能忘记用 cudaFree 将分配的 GPU 显存给释放。

1
2
3
4
5
6
// Copy output from device
unsigned char* image_y = new unsigned char[width*height];
cudaMemcpy(image_y, dev_output, mem_size, cudaMemcpyDeviceToHost);

cudaFree(dev_input);
cudaFree(dev_output);

Timing

使用 CUDA C++ 的接口可以方便的计算运行时间,无论是否是 CUDA Kernel 代码:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
#include "cuda_runtime.h"

// Timing
cudaEvent_t start, stop;
cudaEventCreate(&start);
cudaEventCreate(&stop);

cudaEventRecord(start);

//
// DO SOMETHING
// WITH OR WITHOUT CUDA
//

// Get elapsed time in ms
cudaEventRecord(stop);

cudaEventSynchronize(stop);
float milliseconds = 0;
cudaEventElapsedTime(&milliseconds, start, stop);

编译

编译需要用到 CUDA 库中的 NVIDIA CUDA Compiler (NVCC),在编译的时候注意 GPU 的架构,需要指定 -arch=compute_35 参数,具体的列表可以在 GPU Feature List 里查看。
下面我写了一个非常简单的 makefile 文件,写了一个 rgb2gray_cuda 的规则来使用系统中的 CUDA 10 下的 nvcc 来编译我的 rgb2gray.cu 源码。

1
2
3
4
5
6
7
8
NVCC ?= /usr/local/cuda-10.1/bin/nvcc
NVCCFLAGS ?= -arch=compute_35

rgb2gray_cuda:
$(NVCC) $(NVCCFLAGS) -o rgb2gray_cuda.out rgb2gray.cu lodepng.cpp

clean:
rm -f *.out

编译时只需要执行 $ make rgb2gray_cuda,即可生成 rgb2gray_cuda.out 可执行文件。可以通过 $ make clean 清除所有编译产生的 .out 文件。

Histogram 与 AtomicAdd

在做直方图统计的时候,多个线程会同时访问直方图结果的数列,为了避免 race condition,需要用到 CUDA 的 atomicAdd 函数:

1
2
3
4
5
6
7
8
9
10
11
12
13
// CUDA Kernel
__global__ void HistogramKernel(unsigned char* d_image, unsigned int* d_histo, int width, int height) {

// input_image size: width*height

int row = blockIdx.y * blockDim.y + threadIdx.y;
int col = blockIdx.x * blockDim.x + threadIdx.x;

if ((row < height) && (col < width)) {
int id = row*width + col;
atomicAdd(&d_histo[d_image[id]], 1);
}
}

进一步的,可以使用 shared memory 来提升性能。Using Shared Memory in CUDA C/C++ 中给出了一些介绍及基本的例子。

附上完整 rgb2gray.cu 代码:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
/* 
* Convert RGB image to grayscale
*
* Load image with lodepng (https://github.com/lvandeve/lodepng)
*
* Sample image: lena.png
*
* Y = 0.2126*R + 0.7152*G + 0.0722*B
*
* */

#include <iostream>
#include <string>

#include "cuda_runtime.h"
#include "lodepng.h"

#define CHANNELS 4 // RGBA in PNG
#define BLOCK_SIZE 16 // Thread block size


// CUDA Kernel
__global__ void RGB2GrayKernel(unsigned char* input_image,
unsigned char* output_image,
int width, int height) {

// input_image size: width*height*Channels
// output_image size: width*height

int row = blockIdx.y * blockDim.y + threadIdx.y;
int col = blockIdx.x * blockDim.x + threadIdx.x;

if ((row < height) && (col < width)) {
// Offset in Grayscale image
int offset = row*width + col;

// Get RGB values
unsigned char r = input_image[offset * CHANNELS + 0];
unsigned char g = input_image[offset * CHANNELS + 1];
unsigned char b = input_image[offset * CHANNELS + 2];

// Convert to grayscale
output_image[offset] = 0.2126f*r + 0.7152f*g + 0.0722f*b;
}
}


/* Main */
int main(int argc, char* argv[]) {
const char* input_file = argc > 1 ? argv[1] : "lena.png";

// Variables
std::vector<unsigned char> image;
unsigned int width, height;

// Load image
unsigned error = lodepng::decode(image, width, height, input_file);

std::cout << width << " x " << height << std::endl;

// Allocate memory for CUDA device
size_t mem_size = width * height * sizeof(unsigned char);
unsigned char* dev_input, *dev_output;

cudaMalloc(&dev_input, mem_size * CHANNELS);
cudaMemcpy(dev_input, image.data(), mem_size * CHANNELS, cudaMemcpyHostToDevice);

cudaMalloc(&dev_output, mem_size);

// Invoke CUDA kernel
dim3 dimBlock(BLOCK_SIZE, BLOCK_SIZE);
dim3 dimGrid(width / dimBlock.x, height / dimBlock.y);

RGB2GrayKernel<<<dimGrid, dimBlock>>>(dev_input, dev_output,
width, height);

// Copy output from device
unsigned char* image_y = new unsigned char[width*height];
cudaMemcpy(image_y, dev_output, mem_size, cudaMemcpyDeviceToHost);

cudaFree(dev_input);
cudaFree(dev_output);

// Prepare for output
std::vector<unsigned char> out_image(image.size(), 255);
for (size_t i = 0; i < width*height; i++) {
size_t offset = i * CHANNELS;
out_image[offset + 0] = out_image[offset + 1] = out_image[offset + 2] = image_y[i];
}

// Save processed image
std::string output_file(input_file);
output_file += ".gray.png";

error = lodepng::encode(output_file, out_image, width, height);

delete[] image_y;
return 0;
}

部分参考链接: