Yolov5 は分類モデルのトレーニングを実行します

データセットの収集

データセット
電車
クラス1_imgs
クラス2_imgs
クラス3_imgs
テスト
クラス1_imgs
クラス2_imgs
クラス3_imgs
有効
クラス1_imgs
クラス2_imgs
クラス3_imgs

データセットは、train:test:valid = 7:2:1 に分割することをお勧めします。

トレーニングモデル

yolov5 環境をダウンロードして構成する

pytorch GPU バージョンで実行する必要があります

git clone https://github.com/ultralytics/yolov5 
cd yolov5
pip install -r requirements.txt  # 下载所需包

トレーニングモデル

データセットをyolo5ディレクトリに置くだけです

train.pyを開き、以下の4つのパラメータを設定します。

parser.add_argument('--model', type=str, default='yolov5s-cls.pt', help='initial weights path')
    parser.add_argument('--data', type=str, default='imagenette160', help='cifar10, cifar100, mnist, imagenet, ...')
    parser.add_argument('--epochs', type=int, default=10, help='total training epochs')
    parser.add_argument('--batch-size', type=int, default=64, help='total batch size for all GPUs')
    parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=224, help='train, val image size (pixels)')

–data はデータセットのパス
–batch-size は各トレーニング バッチ内の画像の数です。この値を増やすとトレーニング速度が向上し、メモリ使用量も増加します
–epochs トレーニング時間

モデルはデフォルトで yolov5s-cls.pt を使用します。各事前トレーニング モデルのパラメーターは次のとおりです。

次に、train.py を実行してモデルをトレーニングします。

使用モデル

from utils.augmentations import classify_transforms
from utils.dataloaders import LoadImages
from utils.general import Profile, check_img_size, non_max_suppression, scale_boxes
from models.common import DetectMultiBackend
from utils.torch_utils import select_device

import pathlib
temp = pathlib.PosixPath
pathlib.PosixPath = pathlib.WindowsPath

import torch
import torch.nn.functional as F


def predict_oneImg(model, dir_path, imgsz = (224, 224), dt = (Profile(), Profile(), Profile())):
    dataset = LoadImages(dir_path, img_size=224,
                         transforms=classify_transforms(imgsz[0]))
    for path, im, im0s, vid_cap, s in dataset:
        with dt[0]:
            im = torch.Tensor(im).to(model.device)
            if len(im.shape) == 3:
                im = im[None]  # expand for batch dim

        # Inference
        with dt[1]:
            results = model(im)

        # Post-process
        with dt[2]:
            pred = F.softmax(results, dim=1)

        classes_names = model.names
        prob_list = pred.tolist()[0]
        top3pre = pred.argsort(descending=True).tolist()[0][:3]

        for i in top3pre:
            print("{}: {:.3}%".format(classes_names[i], prob_list[i] * 100), end=" ")
        print('')


device = ''
device = select_device(device)
# 模型路径
model_cls_o = DetectMultiBackend('weights/ball_card_cls2optic.pt', device=device)
# 图片路径
predict_oneImg(model_cls_o, dir_path=r"D:\Code\ML\images\Mywork2\Url_O\16-17")

おすすめ

転載: blog.csdn.net/YierAnla/article/details/128286846