Link

실습: GPU 사용하기

이 책을 읽는 분들의 대부분이 알고계시겠지만, 딥러닝의 연산량은 엄청나기 때문에, 대부분의 딥러닝 연산은 보통 엔비디아Nvidia 그래픽처리장치(GPU)에서 이루어지게 됩니다. 엔비디아에서는 일찌감치 GPU를 통한 병렬연산의 수요에 대해서 눈치채고, 딥러닝이 유행하기 이전부터 엔비디아 GPU를 통한 프로그래밍이 가능하도록 CUDA[1]를 지원해왔습니다. 파이토치에서도 CUDA를 통한 GPU 연산을 지원합니다. 앞서 구현된 파이토치 코드들은 모두 CPU에서 연산이 진행되도록 되어 있었습니다. 이 코드들이 GPU에서 동작하도록 하는 방법들을 알아보도록 하겠습니다.

[1]: CUDACompute Unified Device Architecture는 엔비디아 GPU에서 수행하는 병렬 연산을 프로그래밍 언어를 통해 구현할 수 있도록 하는 기술입니다.

cuda 함수

앞서 우리는 원하는 크기의 임의을 값을 가진 텐서를 생성하는 방법을 배웠습니다. 다음의 코드와 같이 비슷한 방법으로 GPU의 메모리 상에 텐서를 생성할 수 있습니다.

>>> x = torch.cuda.FloatTensor(2, 2)
>>> x
tensor([[0., 0.],
        [0., 0.]], device='cuda:0')

텐서 x의 device가 cuda:0 으로 잡혀있는 것을 볼 수 있습니다. cuda: 뒤에 붙는 숫자는 GPU 디바이스의 인덱스index를 의미합니다. 즉, 첫 번째 디바이스인 0번 GPU를 의미합니다.

앞서의 방법 이외에도, 텐서의 cuda 함수를 통해, CPU 메모리 상에 선언된 텐서를 GPU로 복사하는 방법도 존재합니다.

>>> x = torch.FloatTensor(2, 2)
>>> x
tensor([[-4.4256e-10,  4.5685e-41],
        [-4.4256e-10,  4.5685e-41]])

>>> x = x.cuda()        
>>> x
tensor([[-4.4256e-10,  4.5685e-41],
        [-4.4256e-10,  4.5685e-41]], device='cuda:0')

이 cuda 함수의 인자에 복사하고자하는 목적지 GPU 장치의 인덱스를 넣어, 원하는 디바이스로 복사할 수도 있습니다.

>>> x = x.cuda(device=1) # If you don't have 2nd gpu, error will be occurred.
>>> x
tensor([[-4.4256e-10,  4.5685e-41],
        [-4.4256e-10,  4.5685e-41]], device='cuda:1')

이 cuda 함수는 텐서 뿐만 아니라, nn.Module의 하위 클래스 객체에도 똑같이 적용할 수 있습니다.

>>> import torch.nn as nn
>>> layer = nn.Linear(2, 2)
>>> layer.cuda(0)

주의할 점은 텐서는 cuda 함수를 통해 원하는 디바이스로 복사가 수행되지만, nn.Module의 하위 클래스 객체의 경우에는 복사가 아닌 이동move이 수행된다는 점입니다.

서로 다른 장치간 연산

서로 다른 장치에 올라가 있는 텐서 또는 nn.Module의 하위 클래스 객체끼리는 연산이 불가합니다. CPU와 GPU에 위치한 텐서들끼리 연산이 불가능 할 뿐만 아니라, 0번 GPU와 1번 GPU 사이의 연산도 불가능합니다.

>>> x = torch.FloatTensor(2, 2)
>>> x
tensor([[8.9683e-44, 0.0000e+00],
        [1.1210e-43, 0.0000e+00]])

>>> x + x.cuda(0)
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

cpu 함수

반대로 필요에 따라 GPU 메모리 상에 있는 텐서를 CPU 메모리로 복사해야 할 수 있습니다. 이때는 cpu 함수를 사용하면 됩니다.

>>> x = torch.cuda.FloatTensor(2, 2)
>>> x
tensor([[0., 0.],
        [0., 0.]], device='cuda:0')

>>> x = x.cpu()
>>> x
tensor([[0., 0.],
        [0., 0.]])

to 함수

파이토치는 원래 cuda 함수와 cpu 함수만 제공했지만, to 함수도 함께 제공하고 있습니다. to 함수는 인자로 원하는 디바이스device에 대한 정보를 담은 객체를 받아, 함수 자신을 호출한 객체를 해당 디바이스로 복사(이동) 시킵니다. 디바이스 정보를 담은 객체는 torch.device 를 통해 생성할 수 있습니다.

>>> cpu_device = torch.device('cpu')
>>> gpu_device = torch.device('cuda:0')

앞서 만든 각 디바이스 객체를 통해 원하는 장치로 복사해봅니다.

>>> x = torch.FloatTensor(2, 2)
>>> x
tensor([[1.1874e+26, 4.5598e-41],
        [1.1973e+26, 4.5598e-41]])

>>> x = x.to(gpu_device)
>>> x
tensor([[1.1874e+26, 4.5598e-41],
        [1.1973e+26, 4.5598e-41]], device='cuda:0')

>>> x = x.to(cpu_device)
>>> x
tensor([[1.1874e+26, 4.5598e-41],
        [1.1973e+26, 4.5598e-41]])

device 속성

텐서는 device 속성을 가지고 있어, 우리가 해당 텐서가 위치한 디바이스를 쉽게 파악할 수 있습니다.

>>> x = torch.cuda.FloatTensor(2, 2)
>>> x.device
device(type='cuda', index=0)

재미있는 점은 nn.Module의 하위 클래스 객체는 해당 속성을 갖고 있지 않다는 점입니다. 따라서 만약 모델이 어느 장치에 올라가있는지 알고 싶다면 다음과 같은 방법을 취할 수 있습니다.

>>> layer = nn.Linear(2, 2)
>>> next(layer.parameters()).device
device(type='cpu')

parameters 함수를 통해 모델 내의 파라미터에 대한 이터레이터iterator를 얻은 후, 첫 번째 파라미터 텐서의 device 속성에 접근합니다. 물론 이 방법은 모델 내부의 파라미터 전체가 같은 디바이스에 위치한다는 가정이 존재해야 하지만, 대부분의 경우에 이 가정은 성립할 것입니다.