PyTorch Hub에서 YOLOv5 로드하기#

📚 본 가이드는 PyTorch Hub(https://pytorch.org/hub/ultralytics_yolov5)에서 YOLOv5 🚀를 로드하는 방법을 설명합니다.

시작하기 전에#

Python>=3.8.0 환경에서 PyTorch>=1.8을 포함하여 requirements.txt를 설치합니다. 모델과 데이터셋은 최신 YOLOv5 릴리스에서 자동으로 다운로드됩니다.

pip install -r https://raw.githubusercontent.com/ultralytics/yolov5/master/requirements.txt

팁

ultralytics/yolov5을 복제할 필요가 없습니다. PyTorch Hub가 코드를 자동으로 가져옵니다.

PyTorch Hub로 YOLOv5 로드하기#

간단한 예제#

이 예제는 PyTorch Hub에서 사전 학습된 YOLOv5s 모델을 model(으)로 로드하고 추론을 위해 이미지를 전달합니다. 'yolov5s'은 가장 가볍고 빠른 YOLOv5 모델입니다. 사용 가능한 모든 모델에 대한 자세한 내용은 README를 참조하세요.

import torch

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

# Image
im = "https://ultralytics.com/images/zidane.jpg"

# Inference
results = model(im)

results.pandas().xyxy[0]
#      xmin    ymin    xmax   ymax  confidence  class    name
# 0  749.50   43.50  1148.0  704.5    0.874023      0  person
# 1  433.50  433.50   517.5  714.5    0.687988     27     tie
# 2  114.75  195.75  1095.0  708.0    0.624512      0  person
# 3  986.00  304.00  1028.0  420.0    0.286865     27     tie

상세 예제#

이 예제는 PIL 및 OpenCV 이미지 소스를 사용한 배치 추론을 보여줍니다. results은 콘솔에 출력하고, runs/hub(으)로 저장하고, 지원되는 환경에서 화면에 표시하고, 텐서 또는 pandas 데이터프레임으로 반환할 수 있습니다.

import cv2
import torch
from PIL import Image

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

# Images
for f in "zidane.jpg", "bus.jpg":
    torch.hub.download_url_to_file("https://ultralytics.com/images/" + f, f)  # download 2 images
im1 = Image.open("zidane.jpg")  # PIL image
im2 = cv2.imread("bus.jpg")[..., ::-1]  # OpenCV image (BGR to RGB)

# Inference
results = model([im1, im2], size=640)  # batch of images

# Results
results.print()
results.save()  # or .show()

results.xyxy[0]  # im1 predictions (tensor)
results.pandas().xyxy[0]  # im1 predictions (pandas)
#      xmin    ymin    xmax   ymax  confidence  class    name
# 0  749.50   43.50  1148.0  704.5    0.874023      0  person
# 1  433.50  433.50   517.5  714.5    0.687988     27     tie
# 2  114.75  195.75  1095.0  708.0    0.624512      0  person
# 3  986.00  304.00  1028.0  420.0    0.286865     27     tie

모든 추론 옵션은 YOLOv5 AutoShape() 순방향 메서드를 참조하세요.

추론 설정#

YOLOv5 모델에는 신뢰도 임계값(confidence threshold), IoU 임계값 등 다양한 추론 속성이 포함되어 있으며, 다음과 같이 설정할 수 있습니다:

model.conf = 0.25  # NMS confidence threshold
model.iou = 0.45  # NMS IoU threshold
model.agnostic = False  # NMS class-agnostic
model.multi_label = False  # NMS multiple labels per box
model.classes = None  # (optional list) filter by class, i.e. = [0, 15, 16] for COCO persons, cats and dogs
model.max_det = 1000  # maximum number of detections per image
model.amp = False  # Automatic Mixed Precision (AMP) inference

results = model(im, size=320)  # custom inference size

장치#

모델 생성 후 원하는 장치(device)로 전송할 수 있습니다:

model.cpu()  # CPU
model.cuda()  # GPU
model.to(device)  # i.e. device=torch.device(0)

모델은 모든 device에서 직접 생성할 수도 있습니다:

model = torch.hub.load("ultralytics/yolov5", "yolov5s", device="cpu")  # load on CPU

팁

입력 이미지는 추론 전에 자동으로 올바른 모델 장치로 전송됩니다.

출력 생략#

_verbose=False을 사용하여 모델을 조용히 로드할 수 있습니다:

model = torch.hub.load("ultralytics/yolov5", "yolov5s", _verbose=False)  # load silently

입력 채널#

기본값인 3 채널 대신 4개의 입력 채널을 가진 사전 훈련된 YOLOv5s 모델을 로드하려면 다음과 같이 합니다:

model = torch.hub.load("ultralytics/yolov5", "yolov5s", channels=4)

이 경우 모델은 사전 훈련된 입력 레이어와 모양이 더 이상 같지 않은 첫 번째 입력 레이어를 제외하고는 사전 훈련된 가중치로 구성됩니다. 입력 레이어는 무작위 가중치로 초기화된 상태로 유지됩니다.

클래스 수#

기본값인 80개 대신 10개의 출력 클래스를 가진 사전 훈련된 YOLOv5s 모델을 로드하려면 다음과 같이 합니다:

model = torch.hub.load("ultralytics/yolov5", "yolov5s", classes=10)

이 경우 모델은 사전 훈련된 출력 레이어와 모양이 더 이상 같지 않은 출력 레이어를 제외하고는 사전 훈련된 가중치로 구성됩니다. 출력 레이어는 무작위 가중치로 초기화된 상태로 유지됩니다.

강제 새로고침(Force Reload)#

위 단계에서 문제가 발생하는 경우, force_reload=True을 설정하면 기존 캐시를 폐기하고 PyTorch Hub에서 최신 YOLOv5 버전의 새로운 다운로드를 강제하는 데 도움이 될 수 있습니다. 캐시된 복사본은 ~/.cache/torch/hub에 있습니다. 해당 폴더를 삭제하면 동일한 효과를 얻을 수 있습니다.

model = torch.hub.load("ultralytics/yolov5", "yolov5s", force_reload=True)  # force reload

화면 캡처 추론#

데스크톱 화면에서 추론을 실행하려면 다음과 같이 합니다:

import torch
from PIL import ImageGrab

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

# Image
im = ImageGrab.grab()  # take a screenshot

# Inference
results = model(im)

멀티 GPU 추론#

YOLOv5 모델은 스레드 추론을 통해 여러 GPU에 병렬로 로드할 수 있습니다:

import threading

import torch

def run(model, im):
    """Performs inference on an image using a given model and saves the output; model must support `.save()` method."""
    results = model(im)
    results.save()

# Models
model0 = torch.hub.load("ultralytics/yolov5", "yolov5s", device=0)
model1 = torch.hub.load("ultralytics/yolov5", "yolov5s", device=1)

# Inference
threading.Thread(target=run, args=[model0, "https://ultralytics.com/images/zidane.jpg"], daemon=True).start()
threading.Thread(target=run, args=[model1, "https://ultralytics.com/images/bus.jpg"], daemon=True).start()

학습#

추론이 아닌 학습용으로 YOLOv5 모델을 로드하려면 autoshape=False을 설정합니다. (처음부터 학습하기 위해) 무작위로 초기화된 가중치로 모델을 로드하려면 pretrained=False을 사용합니다. 이 경우 사용자가 직접 학습 스크립트를 제공해야 합니다. 또는 모델 학습에 대해서는 YOLOv5 커스텀 데이터 학습 튜토리얼을 참조하세요.

import torch

model = torch.hub.load("ultralytics/yolov5", "yolov5s", autoshape=False)  # load pretrained
model = torch.hub.load("ultralytics/yolov5", "yolov5s", autoshape=False, pretrained=False)  # load scratch

Base64 결과#

API 서비스와 함께 사용할 때 사용합니다. 자세한 내용은 Flask REST API 예제를 참조하세요.

import base64
from io import BytesIO

from PIL import Image

results = model(im)  # inference

results.ims  # array of original images (as np array) passed to model for inference
results.render()  # updates results.ims with boxes and labels
for im in results.ims:
    buffered = BytesIO()
    im_base64 = Image.fromarray(im)
    im_base64.save(buffered, format="JPEG")
    print(base64.b64encode(buffered.getvalue()).decode("utf-8"))  # base64 encoded image with results

크롭된 결과#

결과를 탐지 영역을 잘라내어(crop) 반환하고 저장할 수 있습니다:

results = model(im)  # inference
crops = results.crop(save=True)  # cropped detections dictionary

Pandas 결과#

결과는 Pandas DataFrames 형식으로 반환될 수 있습니다:

results = model(im)  # inference
results.pandas().xyxy[0]  # Pandas DataFrame

Pandas Output (click to expand)

print(results.pandas().xyxy[0])
#      xmin    ymin    xmax   ymax  confidence  class    name
# 0  749.50   43.50  1148.0  704.5    0.874023      0  person
# 1  433.50  433.50   517.5  714.5    0.687988     27     tie
# 2  114.75  195.75  1095.0  708.0    0.624512      0  person
# 3  986.00  304.00  1028.0  420.0    0.286865     27     tie

정렬된 결과#

결과는 열별로 정렬할 수 있습니다. 예를 들어, 자동차 번호판 숫자 탐지 결과를 왼쪽에서 오른쪽(x축)으로 정렬하려면 다음과 같이 합니다:

results = model(im)  # inference
results.pandas().xyxy[0].sort_values("xmin")  # sorted left-right

JSON 결과#

.to_json() 메서드를 사용하여 .pandas() 데이터프레임으로 변환되면 결과를 JSON 형식으로 반환할 수 있습니다. JSON 형식은 orient 인수를 사용하여 수정할 수 있습니다. 자세한 내용은 pandas .to_json() 문서를 참조하세요.

results = model(ims)  # inference
results.pandas().xyxy[0].to_json(orient="records")  # JSON img1 predictions

JSON Output (click to expand)

[
    {
        "xmin": 749.5,
        "ymin": 43.5,
        "xmax": 1148.0,
        "ymax": 704.5,
        "confidence": 0.8740234375,
        "class": 0,
        "name": "person"
    },
    {
        "xmin": 433.5,
        "ymin": 433.5,
        "xmax": 517.5,
        "ymax": 714.5,
        "confidence": 0.6879882812,
        "class": 27,
        "name": "tie"
    },
    {
        "xmin": 115.25,
        "ymin": 195.75,
        "xmax": 1096.0,
        "ymax": 708.0,
        "confidence": 0.6254882812,
        "class": 0,
        "name": "person"
    },
    {
        "xmin": 986.0,
        "ymin": 304.0,
        "xmax": 1028.0,
        "ymax": 420.0,
        "confidence": 0.2873535156,
        "class": 27,
        "name": "tie"
    }
]

사용자 지정 모델#

이 예제는 PyTorch Hub를 사용하여 커스텀 20개 클래스 VOC로 학습된 YOLOv5s 모델 'best.pt'을 로드합니다.

import torch

model = torch.hub.load("ultralytics/yolov5", "custom", path="path/to/best.pt")  # local model
model = torch.hub.load("path/to/yolov5", "custom", path="path/to/best.pt", source="local")  # local repo

TensorRT, ONNX 및 OpenVINO 모델#

PyTorch Hub는 커스텀 학습된 모델을 포함하여 대부분의 YOLOv5 내보내기 형식에 대한 추론을 지원합니다. 모델 내보내기에 대한 자세한 내용은 TFLite, ONNX, CoreML, TensorRT 내보내기 튜토리얼을 참조하세요.

팁

GPU 벤치마크에서 TensorRT는 PyTorch보다 2~5배 더 빠를 수 있습니다.
CPU 벤치마크에서 ONNX 및 OpenVINO는 PyTorch보다 2~3배 더 빠를 수 있습니다.

import torch

model = torch.hub.load("ultralytics/yolov5", "custom", path="yolov5s.pt")  # PyTorch
model = torch.hub.load("ultralytics/yolov5", "custom", path="yolov5s.torchscript")  # TorchScript
model = torch.hub.load("ultralytics/yolov5", "custom", path="yolov5s.onnx")  # ONNX
model = torch.hub.load("ultralytics/yolov5", "custom", path="yolov5s_openvino_model/")  # OpenVINO
model = torch.hub.load("ultralytics/yolov5", "custom", path="yolov5s.engine")  # TensorRT
model = torch.hub.load("ultralytics/yolov5", "custom", path="yolov5s.mlmodel")  # CoreML (macOS-only)
model = torch.hub.load("ultralytics/yolov5", "custom", path="yolov5s.tflite")  # TFLite
model = torch.hub.load("ultralytics/yolov5", "custom", path="yolov5s_paddle_model/")  # PaddlePaddle

지원되는 환경#

Ultralytics는 프로젝트를 빠르게 시작할 수 있도록 CUDA, CUDNN, Python, PyTorch와 같은 필수 종속성이 미리 설치된 다양한 즉시 사용 가능한 환경을 제공합니다.

무료 GPU 노트북:
Google Cloud: GCP 퀵스타트 가이드
Amazon: AWS 빠른 시작 가이드
Azure: AzureML 빠른 시작 가이드
Docker: Docker 퀵스타트 가이드

프로젝트 상태#

이 배지는 모든 YOLOv5 GitHub Actions CI(Continuous Integration) 테스트가 성공적으로 통과되었음을 나타냅니다. 이 CI 테스트들은 학습, 검증, 추론, 내보내기, 벤치마크 등 다양한 핵심 측면에서 YOLOv5의 기능과 성능을 엄격하게 확인합니다. 테스트는 24시간마다 그리고 새로운 커밋이 발생할 때마다 수행되며 macOS, Windows, Ubuntu에서 일관되고 안정적인 작동을 보장합니다.

기여자

GLglenn-jocher¹

생성됨 2개월 전업데이트됨 2개월 전