CUDA 커널을 Python에 연결하기

2022년 11월 18일 · 원문 · 한국어 번역(원문에서 변경됨)

가우시안 블러는 가우시안 커널을 사용하여 이미지를 부드럽게 만드는 일반적인 이미지 처리 연산입니다. CPU에서 구현하기는 간단하지만, 큰 이미지에서는 상당히 느릴 수 있습니다. CUDA를 사용하여 이를 효율적으로 구현하는 방법을 살펴보고, 쉽게 사용할 수 있도록 Python에 연결해 봅시다.

먼저 NVIDIA CUDA 컴파일러(nvcc)가 포함된 CUDA 툴킷을 설치하십시오:

# Add NVIDIA package repositories
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.0-1_all.deb
sudo dpkg -i cuda-keyring_1.0-1_all.deb
sudo apt-get update
sudo apt-get install build-essential

# Install CUDA toolkit
sudo apt-get -y install cuda-toolkit-12-x

설치를 확인하십시오:

nvcc --version

Python 가상 환경을 만들고 필요한 패키지를 설치하십시오:

pip install numpy matplotlib pillow pycuda

모든 것이 제대로 작동하는지 확인하려면 다음의 간단한 테스트를 실행하십시오:

import pycuda.autoinit
import pycuda.driver as cuda
print(f"CUDA device name: {cuda.Device(0).name()}")

오류 없이 GPU 이름이 출력되면 준비가 완료된 것입니다. 먼저 CUDA C로 블러 연산을 구현하는 방법을 살펴봅시다. 우리 구현의 Core는 각 픽셀을 병렬로 처리하는 커널입니다:

__global__ void gaussianBlur(
    const unsigned char* input,
    unsigned char* output,
    int width,
    int height,
    float sigma
) {
    const int x = blockIdx.x * blockDim.x + threadIdx.x;
    const int y = blockIdx.y * blockDim.y + threadIdx.y;

    if (x >= width || y >= height) return;

    const int kernel_size = 2 * (int)(3.0f * sigma) + 1;
    const int radius = kernel_size / 2;

    float sum = 0.0f;
    float weightSum = 0.0f;

    for (int ky = -radius; ky <= radius; ky++) {
        for (int kx = -radius; kx <= radius; kx++) {
            const int px = min(max(x + kx, 0), width - 1);
            const int py = min(max(y + ky, 0), height - 1);

            const float distance = (kx * kx + ky * ky);
            const float weight = expf(-distance / (2.0f * sigma * sigma));

            sum += input[py * width + px] * weight;
            weightSum += weight;
        }
    }

    output[y * width + x] = (unsigned char)(sum / weightSum);
}

이 코드를 gaussian_blur.cu라는 파일에 저장하십시오.

이 커널은 각 CUDA 스레드가 출력 이미지의 픽셀 하나를 처리하는 방식으로 작동합니다. 각 픽셀에 대해 가우시안 분포를 사용하여 주변 픽셀의 가중 평균을 계산합니다. 커널의 크기는 시그마에 의해 결정됩니다. 가우시안의 유의미한 값 대부분을 포착하기 위해 $3 σ$ 의 반경을 사용합니다.

각 스레드는 먼저 블록 인덱스와 스레드 인덱스를 사용하여 이미지 내에서 자신의 위치를 계산합니다. 그런 다음 픽셀 주변의 정사각형 이웃 영역을 순회하면서 중심 픽셀로부터의 거리에 따라 가우시안 가중치를 계산합니다. 가중치는 2D 가우시안 함수 $exp (- \frac{r ^{2}}{2 σ ^{2}})$ 를 따르며, 여기서 \(r))은 중심으로부터의 거리입니다.

이미지 경계를 올바르게 처리하기 위해 min 및 max 연산을 사용하여 픽셀 좌표를 클램프합니다. 이렇게 하면 이미지 경계 밖을 읽게 되는 경우 가장자리 픽셀이 사실상 확장됩니다. CUDA 구현은 효율적이지만, Python에서 사용하는 것이 더 편리합니다. PyCUDA를 사용하면 Python 코드에서 CUDA 커널을 컴파일하고 실행할 수 있습니다:

import numpy as np
import pycuda.autoinit
import pycuda.driver as cuda
from pycuda.compiler import SourceModule

def gaussian_blur(image, sigma=1.5):
    # Ensure image is grayscale and float32
    if len(image.shape) > 2:
        image = np.mean(image, axis=2)
    image = image.astype(np.float32)

    # CUDA kernel as a string
    cuda_code = open("gaussian_blur.cu").read()

    # Compile the kernel
    mod = SourceModule(cuda_code)
    gaussian_blur_kernel = mod.get_function("gaussianBlur")

    # Prepare input and output arrays
    height, width = image.shape
    output = np.empty_like(image)

    # Copy input to device
    d_input = cuda.mem_alloc(image.nbytes)
    d_output = cuda.mem_alloc(output.nbytes)
    cuda.memcpy_htod(d_input, image)

    # Set up grid and block dimensions
    block_size = (16, 16, 1)
    grid_size = (
        (width + block_size[0] - 1) // block_size[0],
        (height + block_size[1] - 1) // block_size[1]
    )

    # Run the kernel
    gaussian_blur_kernel(
        d_input,
        d_output,
        np.int32(width),
        np.int32(height),
        np.float32(sigma),
        block=block_size,
        grid=grid_size
    )

    # Copy result back to host
    cuda.memcpy_dtoh(output, d_output)

    return output

Python 래퍼는 모든 CUDA 설정과 메모리 관리를 처리합니다. 먼저 입력 이미지가 올바른 형식(그레이스케일 및 float32)인지 확인합니다. CUDA 커널은 문자열로 저장되며, PyCUDA의 SourceModule을 사용하여 런타임에 컴파일됩니다. 그런 다음 GPU에 입력 이미지와 출력 이미지를 위한 메모리를 할당하고 입력 데이터를 GPU로 복사합니다. 커널은 16x16 스레드 블록 크기로 실행되며, 그리드 크기는 전체 이미지를 포괄하도록 계산됩니다.

커널이 실행된 후 결과를 CPU 메모리로 다시 복사하여 반환합니다. 이 함수는 다음과 같이 사용할 수 있습니다:

import matplotlib.pyplot as plt
from PIL import Image

# Load and convert image to grayscale numpy array
img = np.array(Image.open('lenna.png').convert('L'))

# Apply blur
blurred = gaussian_blur(img, sigma=2.0)

# Display results
plt.figure(figsize=(12, 6))
plt.subplot(121)
plt.imshow(img, cmap='gray')
plt.title('Original')
plt.subplot(122)
plt.imshow(blurred, cmap='gray')
plt.title('Blurred')
plt.show()

이 구현은 CPU 기반 구현보다 훨씬 빠르며, 이미지가 클수록 그 차이가 더 커집니다. CUDA는 병렬로 동작하므로 많은 픽셀을 동시에 처리할 수 있으며, 커널의 효율적인 메모리 접근 패턴도 좋은 성능을 유지하는 데 도움이 됩니다. 하지만 CPU와 GPU 사이에서 메모리를 복사하는 오버헤드는 여전히 남아 있습니다. 처리하기 전에 입력 이미지를 GPU 메모리로 전송해야 하고, 처리가 끝나면 결과를 다시 CPU 메모리로 복사해야 합니다. 실제 애플리케이션에서 최적의 성능을 얻으려면 여러 연산을 수행할 때 데이터를 가능한 한 GPU에 유지하여 이러한 메모리 전송을 최소화하는 것이 좋습니다.

Keyboard shortcuts

Stephen Diehl 블로그 한국어 번역

CUDA 커널을 Python에 연결하기