[Pytorch] 장치간 모델 불러오기 (GPU / CPU)

1. GPU에서 모델을 저장하고 CPU에서 불러오기

2. GPU에서 저장하고 GPU에서 불러오기 

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

 

1. GPU에서 모델을 저장하고 CPU에서 불러오기

GPU 모델 저장

torch.save(model.state_dict(), PATH)

저장한 모델을 CPU에서 로드

device = torch.device('cpu')
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH, map_location=device))

위와 같이 GPU에서 학습한 모델을 CPU에서 불러올 때는 torch.load() 함수의 map_location 인자에 'cpu' 값을 전달함

이 경우에는 Tensor에 저장된 내용들은 map_location 인자를 사용하여 CPU 장치에 동적으로 재배치됨 

 

2. GPU에서 모델을 저장하고 GPU에서 불러오기 

GPU모델 저장 (위와 동일)

torch.save(model.state_dict(), PATH)

저장한 모델을 GPU에서 로드

device = torch.device("cuda")
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.to(device)
# 모델에서 사용하는 input Tensor들은 input = input.to(device) 을 호출해야 합니다.

GPU에서 학습된 모델을 GPU에서 불러올 때에는

초기화된 model에 model.to(torch.device('cuda')) 를 호출하여 CUDA 최적화된 모델로 변환해야함

 

또한, 모델에 데이터를 제공하는 모든 입력에 .to(torch.device('cuda')) 함수를 호출해야함

 

my_tensor.to(device) 를 호출하면 GPU에 my_tensor의 복사본을 반환하기 때문에 Tensor를 직접 덮어써야함

my_tensor = my_tensor.to(torch.device('cuda'))

 

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

CPU모델 저장 (위와 동일)

torch.save(model.state_dict(), PATH)

저장한 모델을 GPU에서 로드

device = torch.device("cuda")
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH, map_location="cuda:0"))  # 사용할 GPU 장치 번호를 선택합니다.
model.to(device)
# 모델에서 사용하는 input Tensor들은 input = input.to(device) 을 호출해야 합니다.

CPU에서 학습한 모델을 GPU에서 불러올 때는 torch.load() 함수의 map_location 인자에 cuda:device_id 을 설정하면 모델이 해당 GPU 장치에 로드됨

 

다음으로 model.to(torch.device('cuda')) 을 호출하여 모델의 매개변수 Tensor들을 CUDA Tensor들로 변환해야 함

 

마지막으로 모든 모델 입력에 .to(torch.device('cuda')) 을 사용하여 CUDA 최적화된 모델을 위한 데이터로 만들어야 함

 

my_tensor.to(device) 를 호출하면 GPU에 my_tensor 의 복사본을 반환하며, 이 동작은 my_tensor 를 덮어쓰지 않기 때문에, Tensor를 직접 덮어써야 함

my_tensor = my_tensor.to(torch.device('cuda')) .

 

torch.nn.DataParallel 모델 저장하기

저장하기:

torch.save(model.module.state_dict(), PATH)

불러오기:

# 사용할 장치에 불러옵니다.

torch.nn.DataParallel 은 병렬 GPU 활용을 가능하게 하는 모델 래퍼(wrapper)입니다. DataParallel 모델을 범용적으로 저장하려면 model.module.state_dict() 을 사용하면 됩니다. 이렇게 하면 원하는 모든 장치에 원하는 방식으로 유연하게 모델을 불러올 수 있습니다.

 

만일 map_location 을 설정하지 않는 경우 아래와 같은 에러를 맞이한다.

 

RuntimeError: Attempting to deserialize object on a CUDA device but torch.cuda.is_available() is False. If you are running on a CPU-only machine, please use torch.load with map_location=torch.device('cpu') to map your storages to the CPU.