파이토치 PyTorch CUDA custom layer 만들기 (3) - setup.py

이전 글:

2021.04.28 - [코딩/PyTorch] - 파이토치 PyTorch CUDA custom layer 만들기 (1) - CUDA 커널

2021.04.30 - [코딩/PyTorch] - 파이토치 PyTorch CUDA custom layer 만들기 (2) - 커널 빌드 및 Python 바인딩

 

이전 포스트에 이어 PyTorch (C++)와 PyTorch (Python)을 연결하는 방법을 정리해야겠다.

 

이 부분은 의외로 간단한것이, setup.py 템플릿을 조금만 손보면

 

PyTorch의 cpp_extension이라는 모듈로 어렵지 않게 해결이 가능하다.

 

 

원리는 생략하고 아래와 같이 setup.py를 작성하자.

 

참고로 나는 setup.py의 위치를 cuda 폴더 바깥으로 했다.

 

전체 폴더 구조는 대충 아래와 같다.

cuda_extension 폴더 아래 setup.py를 놓는다.

# setup.py
from os import path
import setuptools

from torch.utils import cpp_extension

def main() -> None:
    setuptools.setup(
        name='ext_vecadd',
        version='1.0.0',
        author='Your name',
        author_email='your@email.com',
        packages=setuptools.find_packages(),
        ext_modules=[cpp_extension.CppExtension(
            name='vecadd_cuda',
            sources=[path.join('cuda', 'vecadd.cpp')],
            libraries=[
                'vecadd_kernel',
            ],
            library_dirs=[path.join('.', 'cuda')],
            extra_compile_args=['-g', '-fPIC'],
        )],
        cmdclass={'build_ext': cpp_extension.BuildExtension},
    )
    return

if __name__ == '__main__':
    main()

당연한 소리지만 version, author, author_email은 취향대로.

 

 

여기서 중요하게 볼 부분은 setup 함수의 ext_modules 부분이다.

 

이 링크에서 공식 documentation을 보면 CppExtension 말고도 CUDAExtension이라는 함수도 존재하는데,

 

막상 써보니까 잘 동작하지 않았고 돌아돌아 결론을 내린게 CppExtension이었다.

 

CUDAExtension을 쓰면 별도의 Makefile로 shared object (.so)를 빌드하지 않아도 될 것 같은데,

 

잘 안되니까 나중에 다시 해봐야겠다.

 

CppExtension의 name 인수에는 나중에 Python에서 해당 모듈을 불러올때 사용하는 이름을 지정한다.

 

예를 들어, 위의 setup.py로 빌드한 모듈은 아래와 같이 불러올 수 있다.

import torch
import vecadd_cuda

반드시 Line 1이 선행되어야 하는 점에 주의하면, 큰 문제는 없다.

• sources 인수는 빌드할 cpp 파일(들)을 지정한다.
  이미 vecadd.cpp에서 CUDA 커널을 호출하도록 되어있기 때문에, .cu 파일을 포함할 필요는 없다.
• 대신에 우리가 빌드한 .so 파일을 libraries라는 인수에 포함시키는데,
  vecadd_kernel을 지정하면 자동으로 libvecadd_kernel.so가 빌드에 포함된다고 생각하면 된다.
• library_dirs에는 해당 .so 파일이 포함되어 있는 폴더를 지정해주고,
• extra_complie_args는 -g와 -fPIC를 넣어주면 완성.

이제 터미널에서 아래와 같은 커맨드로 전체 프로젝트를 빌드한다.

$ python setup.py install

여기서 주의할 점은, conda 등의 가상 환경을 사용 중인 경우 해당 environment에만 패키지가 설치된다는 것.

 

뭔가 더 해야할 것 같은 느낌이 들지만 의외로 지금 당장 작성한 CUDA 커널을 PyTorch (Python)에서 사용 가능하다.

 

우선 Python을 실행시키고, 아래와 같이 하면 테스트 가능.

>>> import torch
>>> import vecadd_cuda
>>> x = torch.randn(4, 4, device=torch.device('cuda'))
>>> y = torch.randn(4, 4, device=torch.device('cuda'))
>>> z = torch.zeros(4, 4, device=torch.device('cuda'))
>>> print(x)
tensor([[ 0.7651,  0.2747, -0.3275, -1.0251],
        [-0.5340, -0.8416,  0.2952,  0.7431],
        [-0.2698, -0.3798, -0.0225,  0.0245],
        [ 1.2056,  1.9831,  0.1669, -0.2347]], device='cuda:0')
>>> print(y)
tensor([[-1.0483,  0.6131, -0.2479,  0.3100],
        [ 0.6286,  0.4175, -0.4093, -1.5263],
        [ 0.6393,  1.3988, -0.0907,  0.1472],
        [-1.0677,  1.0088, -1.4845, -0.7323]], device='cuda:0')
>>> print(z)
tensor([[0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.]], device='cuda:0')
>>> vecadd_cuda.vecadd(x, y, z)
>>> print(z)
tensor([[-0.2832,  0.8878, -0.5754, -0.7151],
        [ 0.0947, -0.4241, -0.1141, -0.7832],
        [ 0.3695,  1.0189, -0.1133,  0.1716],
        [ 0.1379,  2.9919, -1.3177, -0.9670]], device='cuda:0')
>>> print(x + y)
tensor([[-0.2832,  0.8878, -0.5754, -0.7151],
        [ 0.0947, -0.4241, -0.1141, -0.7832],
        [ 0.3695,  1.0189, -0.1133,  0.1716],
        [ 0.1379,  2.9919, -1.3177, -0.9670]], device='cuda:0')

random 수치들은 테스트할 때마다 다르게 나오겠지만, 아무튼 z = x + y가 잘 작동한다.

 

코드를 잘(?) 만들어놔서인지, 2D 텐서를 처리하는 데에도 큰 문제가 없다.

 

 

그런데 현재 상태로는 CUDA 커널이 수정될때마다 Makefile 후 setup.py를 install 하는 두 단계를 거쳐야 해서 조금 번거롭다.

 

기왕 하는거 두 개를 한 번에 하는 스크립트도 작성해놨다.

 

아래 내용을 setup.py가 있는 폴더에 Makefile 파일을 만들고 붙여넣으면 된다.

SUBDIRS = cuda

all: vecadd
    python setup.py install

vecadd:
    $(MAKE) -C cuda

clean:
    rm -rf build dist ./*.egg-info
    for dir in $(SUBDIRS); do $(MAKE) -C $$dir Makefile $@; done

역시나 야매지만 $(MAKE)와 -C라는 문법을 통해 sub directory의 Makefile을 실행하는 형식인가 보다.

 

clean의 경우 저것 말고도 여러 가지 방법이 있는데 내가 보기에 깔끔한 걸로 골랐다.

 

아무튼 이제 setup.py를 실행시키는 대신에, Makefile을 한 번만 실행하면 된다!

$ make

여기서 끝?

 

이면 참 좋겠지만, 몇 가지 불편한 사실들이 남아있다.

1. z가 자동으로 할당되지 않는다.
   보통 우리는 x와 y를 가지고 있을 때, x + y가 z를 return 하기를 원하지
   z를 미리 할당하고 x + y를 채우는 것을 별로 선호하지 않는다.
2. 예외 처리가 되어있지 않다.
   작성한 커널은 CUDA 텐서에 대해서만 작동하지만 CPU 텐서가 들어와도 아무 경고가 없다 (내가 가장 삽질한 부분).
   x + y가 정상적으로 계산되지 않는 것은 당연하고.
3. backward()가 작동하지 않는다. = backpropagation이 불가능하다.
   사실 가장 중요한 부분이다.
   달리 말하면, custom operation으로의 기능은 하지만 custom layer로의 기능은 전혀 할 수가 없다.

1과 2는 그러려니 해도, 3은 조금 치명적이다.

 

조금 불행한 소식은 3 또한 CUDA 커널을 직접 작성해서 구현해야 한다는 점,

 

그나마 다행인 소식은 CUDA 커널만 추가로 작성하고

 

Makefile이나 vecadd.cpp, setup.py를 크게 고칠 필요는 없다는 점이다.

 

또한 backward가 사용 가능하게 하려면 autograd.Function을 상속받아야 하는데, 이를 통해 1과 2를 함께 해결 가능하다.

 

 

이전에 backward를 미리 작성해놓지 않은 이유는 사실 간단한데,

 

이 단계에서 구현한 연산이 정확히 동작하지 않는다면 backward 또한 제대로 동작 할리 만무하기 때문이다.

 

그러니까 일단 CUDA 커널을 제대로 작성했는지 PyTorch (Python) 인터페이스를 사용하여 편하게 확인하고,

 

문제가 있다면 고쳐서 완벽한 동작을 보장한 다음 나머지 잡일 (좀 많지만)을 처리하러 가면 된다.

 

다음 글:

2021.05.01 - [코딩/PyTorch] - 파이토치 PyTorch CUDA custom layer 만들기 (4) - autograd.Function과 backward 구현

파이토치 PyTorch CUDA custom layer 만들기 (4) - autograd.Function과 backward 구현