파이토치(pytorch)에서 사전학습 모델과 데이터셋의 경로 설정

파이토치를 사용할 때, 사전학습 된 모델을 불러오거나 토치비전 데이터셋(torchvision.datasets)을 불러올 때, 디스크에 용량을 차지하게 된다. 따로 설정하지 않으면 토치에서 기본으로 설정한 디렉토리에 이런 것들이 저장되는 데, 컴퓨터에 추가적인 저장장치를 마운트해서 사용 중이면 경로를 바꿔서 설정할 수 있다.

 

1. 사전학습 된 모델의 경로 설정

TORCH_HOME 이라는 환경 변수에 사전학습 된 모델 파라미터를 저장할 경로를 설정한다.

(1) 어느 코드를 사용하건 변경할 경로를 설정해두고 그대로 쓰려면 아래처럼 사용자의 홈 디렉토리에 있는 .bashrc 와 같은 곳에 export 명령어로 TORCH_HOME 환경 변수를 설정한다.

 

(2) 파이썬 코드에서 마다 디렉토리를 바꿔주려면 아래 처럼 os를 임포트 한 후 내가 사용할 경로를 설정한다.

2. 토치비전 데이터셋 경로 설정

토치비전 데이터셋의 경로를 내 마음대로 설정하려면, 번거롭지만 torchvision.datasets 의 데이터 클래스 마다 root 인자에 특정 디렉토리를 지정해주어야 한다. 위에서 TORCH_HOME  을 지정해주었으니 이를 이용해서 경로를 설정하면, 사전학습 모델의 체크 포인트와 데이터셋 각각을 통일성 있는 경로에 저장할 수 있다. 예시 코드는 아래와 같다.

 

 

 

참고

1. 사전학습 된 모델 경로 설정 https://pytorch.org/vision/stable/models.html

Models and pre-trained weights — Torchvision 0.17 documentation

2. 데이터셋 경로 설정

https://pytorch.org/vision/0.8/datasets.html#cifar

torchvision.datasets — Torchvision 0.8.1 documentation