[PyTorch,YOLO, hub] YOLOv5 torch hub 로 인터넷 연결 없이 Load하기

YOLOv5 모델은 torch.hub.load를 이용하여 Load 가능하다.

 

https://colab.research.google.com/github/ultralytics/yolov5/blob/master/tutorial.ipynb#scrollTo=GMusP4OAxFu6

YOLOv5 Tutorial

# YOLOv5 PyTorch HUB Inference (DetectionModels only)
import torch

model = torch.hub.load('ultralytics/yolov5', 'yolov5s', force_reload=True)  # yolov5n - yolov5x6 or custom
im = 'https://ultralytics.com/images/zidane.jpg'  # file, Path, PIL.Image, OpenCV, nparray, list
results = model(im)  # inference
results.print()  # or .show(), .save(), .crop(), .pandas(), etc.

 

그런데, 이때 모델 Weight이 이전에 다운로드가 되었었더라도 모델이 Load 되려면 인터넷이 연결돼있어야한다.

 

model을 최초 한 번 다운로드 받은 후에, 인터넷 연결 없이도 이미 다운로드된 Weight 파일을 사용하여 Load를 하려면 어떻게 해야할까?

 

일단 온라인 환경에서 위 코드를 실행시키면 코드가 존재하는 경로에 Weight이 저장(아래 사진에서 yolov5s.pt)될 것이다. yolov5_torchload.py 파일에는 위 코드가 작성돼있다.

왜 Weight 파일을 다운로드 받았음에도 인터넷 연결이 되어야 위 코드가 정상적으로 실행될 수 있을까?

 

PyTorch의 Weight파일에는 모델의 Forward 과정이 정의돼있지 않아서 이 Weight 만 있어서는 모델 코드를 Load 할 수 가 없다. 모델의 Forward 과정은 'ultralytics/yolov5' 프로젝트로 부터 모델 코드까지 다운로드 받아야 실행될 수 있다. 그러면 해결 방법은 Weight 뿐만 아니라 내 PC에  'ultralytics/yolov5' 프로젝트도 다운로드를 받아 놓는 것이다.

 

git clone https://github.com/ultralytics/yolov5.git

 

그리고 코드를 아래와 같이 수정하면 오프라인 환경에서 yolov5 모델이 정상적으로 동작하는 것을 확인할 수 있다.

 

# YOLOv5 PyTorch HUB Inference (DetectionModels only)
import torch

model = torch.hub.load('./yolov5', 'custom',  'yolov5s.pt', source='local')  # yolov5n - yolov5x6 or custom
im = './test.png'  # file, Path, PIL.Image, OpenCV, nparray, list
results = model(im)  # inference
results.print()  # or .show(), .save(), .crop(), .pandas(), etc.

 

model = torch.hub.load('./yolov5', 'custom',  'yolov5s.pt', source='local')

 

위 코드에서는 model 정의를 위와 같이 해주었는데, 맨 앞에 경로는 yolov5 가 설치된 경로이면되고, yolov5s.pt도 v5m 을 쓴다거나 다른 경로에 있다하면 그냥 경로 맞춰서 써주면된다. 저 source 인자가 바로 Local에 다운로드된 프로젝트 폴더로 부터 코드를 Load 할 수 있게끔 해주는 인자이다.

 

코드를 실행시켜보면

 

아래 이미지(위 코드에서 test.png) 를 입력하였을때, 마지막 줄 코드를 .print() 대신 .show()로 바꾸면

 

 

아래와 같은 결과를 정성적으로 확인할 수 있을 것이다.

 

 

 

yolov5/hubconf.py 보면 이미 이 사용법이 주석으로 친절히 작성돼있는 걸 확인할 수 있다.

 

https://github.com/ultralytics/yolov5/blob/master/hubconf.py

GitHub - ultralytics/yolov5: YOLOv5 🚀 in PyTorch > ONNX > CoreML > TFLite

 

언젠가 내 모델도 torch hub로 배포해보고 싶다. 23년도에 해봐야지!

 

+) 이와 관련해서 PR 까지 Merge!

 

https://github.com/ultralytics/yolov5/pull/10460

Add force_reload=True when loading model using torch hub by developer0hye · Pull Request #10460 · ultralytics/yolov5