Yolov5 crea un conjunto de datos y entrena un modelo de clasificación
Recopilación de conjuntos de datos
-
conjuntos de datos
-
-
tren
- clases1_imgs
- clases2_imgs
- clases3_imgs
- …
-
-
prueba
- clases1_imgs
- clases2_imgs
- clases3_imgs
- …
-
-
válido
- clases1_imgs
- clases2_imgs
- clases3_imgs
- …
Se recomienda dividir el conjunto de datos en tren: prueba: válido = 7: 2: 1
Modelo de entrenamiento
Descargue y configure el entorno yolov5
Debe realizarse bajo la versión pytorch gpu
git clone https://github.com/ultralytics/yolov5
cd yolov5
pip install -r requirements.txt # 下载所需包
Modelo de entrenamiento
Simplemente coloque el conjunto de datos en el directorio yolo5
Abra train.py y configure los siguientes cuatro parámetros
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 es la ruta del conjunto de datos
–batch-size es el número de imágenes en cada lote de entrenamiento. Incrementar este valor aumentará la velocidad de entrenamiento y también dará como resultado un mayor uso de memoria
–tiempos de entrenamiento de épocas
El modelo utiliza yolov5s-cls.pt de forma predeterminada. Los parámetros de cada modelo previamente entrenado son los siguientes:
Luego ejecuta train.py para entrenar el modelo.
Usar modelo
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")