yolov5 로컬에서 로드하기

토치로 yolov5모델을 로드하려고 하는데

model = torch.hub.load(‘ultralytics/yolov5’, ‘custom’, path=‘D:/best.pt’, force_reload=True)

hub가 아닌 로컬에서 하는 방법이 없을까요?ㅠㅠ

안녕하세요, @1112 님.

글에서 적어주신 코드는 PyTorch Hub를 사용(torch.hub.load())하여 불러오는 것을 말씀해주신 것 같은데요, 맞으실까요?
(yolov5의 아래 예시처럼 사용하시기 위해 가져오신 것처럼 보입니다. :slight_smile: )

import torch

# Model
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')

# Image
img = 'https://ultralytics.com/images/zidane.jpg'

# Inference
results = model(img)

다만 질문해주신 것은 이해가 잘 안되었는데요,
로컬에서 가져오신다는 말씀은 위와 같이 가져온 yolov5 모델을 로컬에 저장하고 불러오는 방법이 궁금하신 것일까요?

혹시 torch.save() 등을 사용하여 일반적인 모델 저장 / 불러오기 방법을 사용하지 못하시는 상황이신지 궁금합니다.

https://tutorials.pytorch.kr/beginner/saving_loading_models.html

감사합니다~

네 맞습니다!! yolov5 모델을 로컬에 저장하고 불러오는 방법이 궁금합니다ㅠㅠ

네, PyTorch Hub에서 불러온 yolov5s 모델을 로컬에 저장한 뒤에,
일반적인 모델을 저장하고 불러오는 방법과 동일하게 모델을 저장하는 방식으로는 원하시는 것을 진행하시기 어려우신지 궁금했습니다. ^^;

일반적인 모델의 저장 / 불러오기는 torch.save() 문서와 아래 튜토리얼을 참고하시면 좋을 것 같습니다.

https://pytorch.org/docs/stable/generated/torch.save.html

https://tutorials.pytorch.kr/beginner/saving_loading_models.html

아…토치에 대한 이해가 없었네요ㅠㅠ 감사합니다~!!!

좋아요 1

네, 혹시나 진행하시면서 다른 궁금하신게 있으시면 알려주세요 :slight_smile:

ModuleNotFoundError: No module named 'models

이라는 이슈가 발생했는데 YOLOv5는 .save를 사용하려면

.state_dict()를 사용해야 하나요??

pkl로 저장을 해야 하나요??

안녕하세요, @1112 님.

남겨주신 에러 메시지는 models라는 모듈이 없다는 것인데요, 어떤 코드를 사용하셨을까요?

저장하는 방식은 모델의 구조(YOLOv5 등)보다는 프레임워크에 따라서 달라지는 것이고,
PyTorch에서는 위 튜토리얼처럼 다양한 방식으로 모델을 저장하고 불러오는 방식을 제공하고 있습니다.

불러오기 단계가 아니라 저장하는 단계에서 위와 같은 에러를 보셨다면
사용하시는 코드의 오류를 먼저 살펴보시는 것이 좋을 것 같습니다. :slight_smile:

import torch

model = torch.hub.load(‘ultralytics/yolov5’, ‘custom’, path=’/content/drive/MyDrive/dataset2/best.pt’, force_reload=True)

yolov5모델을 커스텀 한후에

torch.save(model, ‘./best.pth’)
모델을 저장 후

model = torch.load(’/content/drive/MyDrive/tmp.pth’)
모델을 로드하려고 했을 때 발생한 문제였습니다ㅜㅜ

안녕하세요, @1112 님.
작성해주신 코드에는 models가 보이지 않는데요,
혹시 전체 에러 메시지도 함께 공유해주실 수 있으실까요?