PyTorch 문서) PyTorch Recipes - Saving and loading models across devices in PyTorch

다양한 디바이스에서 모델을 저장하고 불러와 사용하기

• 원글 : Saving and loading models across devices in PyTorch
• 디바이스를 바꿔가면서 여러 모델을 저장하거나 불러와야 할 때가 있다.
  • 참고) 여기서 디바이스(device)는 신경망을 훈련하거나 추론할 때, 연산이 일어나는 장치를 말하며, GPU를 사용한 가속연산 여부에 따라 cpu와 cuda로 구분되고 있다.

Introduction

• 서로 다른 디바이스 사이에서 모델을 저장하고 불러와 사용하는 과정은 파이토치에서는 비교적 직관적이다.
• 이 글에서는 CPU와 GPU를 바꿔가며 모델의 저장과 로드를 확인한다.

Setup

• 구글 코랩 등의 실행환경에서, 이 글에 포함된 GPU 디바이스에서의 코드 실행이 제대로 이뤄지게 하려면,
  • 런타임 유형을 "GPU"나 다른 가속연산장치(예-TPU)로 변경해야한다.

pip install torch

Steps

1. 데이터를 로드하기 위해 필요한 라이브러리를 불러온다.
2. 신경망을 정의하고 초기화한다.
3. GPU에서 저장하고, CPU에서 불러오기
4. GPU에서 저장하고, GPU에서 불러오기
5. CPU에서 저장하고, GPU에서 불러오기
6. 여러 디바이스에 분산된 DataParallel 모델을 저장하고 불러오기

#단계 1. 데이터를 로드하기 위해 필요한 라이브러리 불러오기

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

#단계 2. 신경망을 정의하고 초기화

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

net = Net()

#단계 3. GPU에서 모델을 저장하고, CPU에서 모델 불러오기

• GPU에서 훈련된 모델을 CPU로 불러오려면
  • torch.load() 함수의 map_location 인자에 torch.device('cpu')를 넘겨준다.
• 텐서에 저장된 내용이 map_location에 따라 CPU에 동적으로 재배치된다.

PATH = "model.pt"
# 모델은 GPU에서 훈련되었다고 가정하고, 저장한다.
torch.save(net.state_dict, PATH)

# CPU에서 모델을 불러온다.
device = torch.device('cpu')
model = Net()
model.load_state_dict(torch.load(PATH, map_location=device, weights_only=True)

#단계 4. GPU에서 모델을 저장하고, GPU에서 모델 불러오기

• GPU에서 훈련되고 저장된 모델을 GPU로 다시 불러와 사용할 때는,
  • model.to(torch.device('cuda'))를 사용하여, 초기화된 모델을 CUDA 최적화 모델로 단순히 변환하면 된다.
• 변환된 모델에 입력되는 데이터 역시 .to(torch.device('cuda'))를 사용해 변환하여 사용해야 한다는 점에 유의한다.
  • 아래에서 my_tensor.to(device)를 호출하면, GPU에서의 my_tensor의 복사본을 반환한다.
  • 즉, 원래의 my_tensor를 알아서 덮어쓰는 in-place연산이 아니므로
  • 다시 my_tensor = my_tensor.to(torch.device('cuda'))와 같이 할당해주는 과정이 필요하다.

# 모델의 저장
torch.save(net.state_dict(), PATH)

# GPU에서 학습되고 저장된 모델을 GPU로 로드하기
device = torch.device("cuda")
model = Net()
model.load_state_dict(torch.load(PATH))
model.to(device)

#단계 5. CPU에서 모델을 저장하고, GPU에서 불러오기

• CPU에서 훈련되고 저장된 모델을, GPU로 불러올 때는
  • torch.load() 함수의 map_location 인자에 cuda:device_id 값을 넘겨줘야 한다.
  • 이는 모델을 지정된 GPU 디바이스로 로드하게 해준다.
• 역시 model.to(torch.device('cuda'))를 호출하여 모델의 매개변수 텐서를 CUDA (최적화)텐서로 변환한다.
• 마지막으로 변환된 모델에 입력될 데이터에도 .to(torch.device('cuda'))를 호출하여, 데이터가 CUDA 최적화 모델에서 사용될 수 있게 한다.

# CPU에서 훈련된 모델의 저장
torch.save(net.state_dict(), PATH)

# GPU로 불러오기
device = torch.device("cuda")
model = Net()
model.load_state_dict(torch.load(PATH, map_location="cuda:0"))
model.to(device)

#단계 6. torch.nn.DataParallel 모델의 저장

• torch.nn.DataParallel은 병렬 GPU를 활용할 수 있게 해주는 모델 래퍼(wrapper)함수다.
• DataParallel 모델을 일반적으로 사용될 수 있게 저장하려면, model.module.state_dict()를 저장하면 된다.
  • 이 방식은 모델을 어떤 디바이스에도 불러올 수 있도록 다양한 방법을 제공하는 유연성을 가진다.

# DataParallel 모델의 저장
torch.save(net.module.state_dict(), PATH)