[PyTorch] File I/O, GPUs

우리는 다양한 구조를 사용하기 위해 결과를 저장하려는 경우가 존재한다.

또한 긴 훈련 과정을 진행할 때도 가장 좋은 방법은 중간 체크 포인트를 주기적으로 저장을 해주어야 며칠 동안 계산한 내용을 잃지 않을 수 있다. 따라서 개별 가중치 벡터와 전체 모델을 모두 로드하고 저장하는 방법에 대해 설명한다. 

 

Loading and Saving Tensors

개별 텐서의 경우 load 및 save 함수를 직접 호출하여 각각 읽고 쓸 수 있다.

import torch
from torch import nn
from torch.nn import functional as F

x = torch.arange(4)
print(x)

torch.save(x, 'x-file')
x2 = torch.load("x-file")
print(x2)

---

Loading and Saving Model parameters

개별 텐서를 저장하는 것은 유용하지만 전체 모델을 저장하려면 번거로운 작업이 될 것이다.

이러한 이유로 딥 러닝 프레임 워크는 전체 모델을 저장하고 로드하는 내장 기능을 제공한다.

주의해야 할 점은 전체 모델 자체를 저장하는 것이 아닌 모델의 매개 변수를 저장한다는 것이다.

model.eval()은 모델을 실행할 때 training모드에서 inference 모드로 전환시켜주는 것이다.

training와 inference시에 각각 다르게 동작하는 layer들이 존재한다.

예를 들어 dropout 또는 BatchNorm 같은 경우는 training에서는 동작하지만 inference에서는 동작하지 않는다.

import torch
from torch import nn
from torch.nn import functional as F

class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.hidden = nn.Linear(20, 256)
        self.output = nn.Linear(256, 10)

    def forward(self, x):
        return self.output(F.relu(self.hidden(x)))

net = MLP()
X = torch.randn(size=(2, 20))
Y = net(X)

torch.save(net.state_dict(), 'mlp.params')

clone = MLP()
clone.load_state_dict(torch.load("mlp.params"))
clone.eval()

Y_clone = clone(X)
print(Y_clone == Y)

---

우리는 지난 시간동안 gpu의 빠른 성능 증가를 경험해왔다. 이러한 큰 기회를 사용할 필요성이 있다.

딥 러닝 프레임 워크는 이러한 gpu 사용 기능을 제공한다.

 

Computing Devices

저장 및 계산을 위해 CPU 및 GPU와 같은 장치를 지정할 수 있다.

tensor는 기본적으로 CPU를 사용하여 계산한다.

torch.device('cpu')는 cpu를 사용, torch.cuda.device('cuda')는 gpu를 사용이다.

다수의 gpu가 있는 경우, torch.cuda.device('cuda:{i}') 또는 torch.device(f'cuda:{i}')로 표현해 사용한다.

torch.cuda.device_count()는 gpu의 개수를 count 하는 함수이다.

import torch
from torch import nn

print(torch.device('cpu'), '\n', torch.cuda.device('cuda'), '\n', torch.cuda.device('cuda:1'), '\n')

print(torch.cuda.device_count())

---

Neural Networks and GPUs

설계한 신경망을 gpu 상에서 돌릴 수 있다.

하지만 이 때 사용할 데이터들은 gpu와 같은 공간에 저장되어서 사용되어야만 계산이 가능하다.

서로 다른 gpu 상에서 사용할 경우는 X.cuda(i)를 이용하여 복사할 수 있다.

예로 들면 만약 gpu:0에 있는 X의 데이터를 gpu:1으로 복사할 경우 X = X.cuda(1)과 같은 방법으로 사용할 수 있다.

 

아래의 소스코드는 다음과 같다.

X와 설게한 신경망 모두 torch.device('cuda')를 통해 동일한 gpu 상에 저장시켜놓아 정상적으로 계산이 가능한 것이다.

import torch
from torch import nn

X = torch.ones(2, 3, device=torch.device('cuda'))

net = nn.Sequential(nn.Linear(3, 1))
net = net.to(device=torch.device('cuda'))

print(net(X))

 

ref. Dive into Deep Learning, Aston Zhang and Zachary C. Lipton and Mu Li and Alexander J. Smola, 2020