CUDAを使って2枚の画像の平均をとるプログラム


このエントリーをはてなブックマークに追加

CUDAの勉強がてら書いてみました。2枚の同サイズの画像を読み込んで、平均画像を作成して保存するだけのプログラムです。環境はUbuntu 16.04です。

// 2枚の画像の平均を取るプログラム
// 画像は以下から6000x4000のものを取得
//     https://www.pexels.com/photo/black-plants-photo-943903/
//     https://www.pexels.com/photo/person-riding-on-zip-line-1090551/

#include <chrono>
#include <cstdio>
#include <iostream>
#include <opencv2/core/core.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/highgui/highgui.hpp>

const int THREADS_PER_BLOCK = 16;

// https://developer.nvidia.com/cuda-example のサンプルから引用
static void HandleError( cudaError_t err, const char *file, int line ) {
    if (err != cudaSuccess) {
        printf( "%s in %s at line %d\n", cudaGetErrorString( err ), file, line );
        exit( EXIT_FAILURE );
    }
}
#define HANDLE_ERROR( err ) (HandleError( err, __FILE__, __LINE__ ))

__global__ void use_gpu_sub(unsigned char *a, unsigned char *b, unsigned char *c, int N) {
    int tid = threadIdx.x + blockIdx.x * blockDim.x;
    if (tid < N) {
        c[tid] = (a[tid] + b[tid]) / 2;
    }
}

void use_gpu(unsigned char* a, unsigned char* b, unsigned char* c, int height, int width, int channel) {
    unsigned char *dev_a, *dev_b, *dev_c;
    const int N = height * width * 3;

    // GPUのためのメモリ確保
    HANDLE_ERROR( cudaMalloc( (void**)&dev_a, N * sizeof(unsigned char) ) );
    HANDLE_ERROR( cudaMalloc( (void**)&dev_b, N * sizeof(unsigned char) ) );
    HANDLE_ERROR( cudaMalloc( (void**)&dev_c, N * sizeof(unsigned char) ) );

    // a, bがさす画像データをGPUのメモリにコピー
    HANDLE_ERROR( cudaMemcpy( dev_a, a, N * sizeof(unsigned char),
                            cudaMemcpyHostToDevice ) );
    HANDLE_ERROR( cudaMemcpy( dev_b, b, N * sizeof(unsigned char),
                            cudaMemcpyHostToDevice ) );

    // (N + THREADS_PER_BLOCK - 1)/THREADS_PER_BLOCK個のブロックをGPU上で実行。
    // 各ブロックでは THREADS_PER_BLOCK 個のスレッドが立つ
    use_gpu_sub<<<(N + THREADS_PER_BLOCK - 1)/THREADS_PER_BLOCK, THREADS_PER_BLOCK>>>( dev_a, dev_b, dev_c, N );

    // GPのメモリにある計算結果をCPU側にコピー
    HANDLE_ERROR( cudaMemcpy( c, dev_c, N * sizeof(unsigned char),
                            cudaMemcpyDeviceToHost ) );

    // free the memory allocated on the GPU
    HANDLE_ERROR( cudaFree( dev_a ) );
    HANDLE_ERROR( cudaFree( dev_b ) );
    HANDLE_ERROR( cudaFree( dev_c ) );
}

void use_cpu(unsigned char* a, unsigned char* b, unsigned char* c, int height, int width, int channel) {
    int N = height * width * channel;
    for (int i = 0; i < N; ++i) {
        c[i] = (a[i] + b[i]) / 2;
    }
}

void average_two_images() {
    cv::Mat img1 = cv::imread("adventure-asia-bali-1090551.jpg");
    cv::Mat img2 = cv::imread("beautiful-close-up-color-943903.jpg");

    assert(img1.size() == img2.size());
    assert(img1.type() == img2.type());
    std::cout << "image size: " << img1.size() << std::endl;

    cv::Mat img3_gpu = cv::Mat(img1.size(), img1.type());
    cv::Mat img3_cpu = cv::Mat(img1.size(), img1.type());

    // gpuを使って
    auto s1 = std::chrono::system_clock::now();
    use_gpu(img1.data, img2.data, img3_gpu.data, img1.rows, img1.cols, img1.channels());
    auto e1 = std::chrono::system_clock::now();
    std::cout << "gpu time: " << std::chrono::duration_cast<std::chrono::milliseconds>(e1 - s1).count() << " ms" << std::endl;

    // cpuを使って
    auto s2 = std::chrono::system_clock::now();
    use_cpu(img1.data, img2.data, img3_cpu.data, img1.rows, img1.cols, img1.channels());    
    auto e2 = std::chrono::system_clock::now();
    std::cout << "cpu time: " << std::chrono::duration_cast<std::chrono::milliseconds>(e2 - s2).count() << " ms" << std::endl;

    imwrite("averaged.jpg", img3_gpu);
    
    // 結果が同じことを確認
    assert(cv::countNonZero(img3_gpu != img3_cpu) == 0);
}

int main()
{
    average_two_images();
    return 0;
}

これを以下のコマンドでビルドします。

nvcc -O2 -std=c++11 average_two_images.cu $(pkg-config --cflags --libs opencv)

以下は実行結果です。GPUは転送速度がネックになっていてCPUより遅いです。

image size: [6000 x 4000]
gpu time: 180 ms
cpu time: 51 ms

コード中で定数にしている THREADS_PER_BLOCKの値を変えると処理時間が変わりました。どうやって最適な値を決めればよいのでしょう?

    176ms: 512スレッド
    172ms: 256スレッド
    158ms: 128スレッド
    158ms: 64スレッド
    180ms: 32スレッド
    180ms: 16スレッド