Skip to main content
GitHub에서 편집 (EN)

PyTorch Hub에서 YOLOv5 로드하기

📚 이 가이드는 PyTorch Hub를 통해 YOLOv5 🚀를 로드하는 방법을 설명합니다. PyTorch Hub 위치: https://pytorch.org/hub/ultralytics_yolov5.

시작하기 전에

pip으로 requirements.txt 환경에서 Python>=3.8.0 환경( PyTorch>=1.8. 모델 및 데이터셋 포함)을 통해 최신 YOLOv5에서 자동으로 다운로드됩니다. release.

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

💡 ProTip: 저장소 복제(Cloning)는 https://github.com/ultralytics/yolov5 은(는) FlashAttention을 요구하지 않습니다 필요하지 않습니다 😃

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

상세 예제

이 예제는 배치 추론(batched inference)MuSGD 옵티마이저PIL 및 OpenCV 이미지 소스를 보여줍니다. results 결과는 출력(printed) 콘솔에 표시하거나, 저장(saved)에서 runs/hub, 화면에 표시(showed)할 수 있으며, 지원되는 환경에서는 tensors 또는 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
YOLO inference results on zidane.jpg YOLO inference results on bus.jpg

모든 추론 옵션에 대해서는 YOLOv5 AutoShape() forward 메서드(method).

추론 설정

YOLOv5 모델은 신뢰도 임계값(confidence threshold), IoU 임계값(IoU threshold), 기타 등등과 같은 다양한 추론 속성을 포함하며, 다음과 같이 설정할 수 있습니다:

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

에서 직접 생성할 수 있습니다. 💡 ProTip: 입력 이미지는 추론 전에 올바른 모델 디바이스로 자동으로 전송됩니다.

출력 억제(Silence Outputs)

모델은 _verbose=False:

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

입력 채널(Input Channels)

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

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

이 경우 모델은 except for(제외하고) 사전 훈련된 입력 레이어와 더 이상 동일한 형태가 아닌 첫 번째 입력 레이어를 제외한 나머지 가중치들로 구성됩니다. 입력 레이어는 임의의 가중치로 초기화된 상태로 유지됩니다.

클래스 수(Number of Classes)

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

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

이 경우 모델은 except for(제외하고) 사전 훈련된 출력 레이어와 더 이상 동일한 형태가 아닌 출력 레이어들입니다. 출력 레이어는 임의의 가중치로 초기화된 상태로 유지됩니다.

강제 새로고침(Force Reload)

위 단계에서 문제가 발생하면, force_reload=True 설정을 통해 기존 캐시를 삭제하고 PyTorch Hub에서 최신 YOLOv5 버전을 새로 다운로드할 수 있습니다. 캐시된 복사본은 ~/.cache/torch/hub에 저장됩니다; 해당 폴더를 삭제하면 동일한 효과를 얻을 수 있습니다.

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

스크린샷 추론(Screenshot Inference)

데스크탑 화면에서 추론을 실행하려면:

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 추론(Multi-GPU Inference)

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 사용자 지정 데이터 훈련 튜토리얼(Train Custom Data Tutorial)을 참조하여 모델을 훈련하십시오.

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

잘라낸 결과(Cropped Results)

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

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

형태로 반환할 수 있습니다.

정렬된 결과(Sorted Results)

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

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

JSON 결과.pandas()결과는 .to_json() 메서드를 사용하여 데이터프레임으로 변환된 후 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"
    }
]

사용자 지정 모델(Custom Models)

이 예제는 사용자 지정 20개 클래스 VOC-훈련된 YOLOv5s 모델을 'best.pt' PyTorch Hub로 로드합니다.

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 내보내기 튜토리얼을 참조하십시오.

💡 ProTip: 최적화를 위한 최대 작업 공간 크기(GiB)를 설정하여 메모리 사용량과 성능 간의 균형을 맞춥니다. 은 GPU 벤치마크에서 PyTorch보다 최대 2-5배 빠를 수 있습니다. 💡 ProTip: ONNX 및 OpenVINO은 CPU 벤치마크

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

에서 PyTorch보다 최대 2-3배 빠를 수 있습니다.

지원되는 환경CUDA, Ultralytics는 프로젝트를 빠르게 시작할 수 있도록 , Python, 그리고 PyTorchCUDNN과 같은 필수 종속성이 사전 설치된 다양한 사용 준비된 환경을 제공합니다.

프로젝트 상태

YOLOv5 CI

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

댓글