You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
81 lines
2.9 KiB
81 lines
2.9 KiB
# Ultralytics YOLO 🚀, AGPL-3.0 license
|
|
|
|
from copy import copy
|
|
|
|
import torch
|
|
|
|
from ultralytics.nn.tasks import RTDETRDetectionModel
|
|
from ultralytics.yolo.utils import DEFAULT_CFG, RANK, colorstr
|
|
from ultralytics.yolo.v8.detect import DetectionTrainer
|
|
|
|
from .val import RTDETRDataset, RTDETRValidator
|
|
|
|
|
|
class RTDETRTrainer(DetectionTrainer):
|
|
|
|
def get_model(self, cfg=None, weights=None, verbose=True):
|
|
"""Return a YOLO detection model."""
|
|
model = RTDETRDetectionModel(cfg, nc=self.data['nc'], verbose=verbose and RANK == -1)
|
|
if weights:
|
|
model.load(weights)
|
|
return model
|
|
|
|
def build_dataset(self, img_path, mode='val', batch=None):
|
|
"""Build RTDETR Dataset
|
|
|
|
Args:
|
|
img_path (str): Path to the folder containing images.
|
|
mode (str): `train` mode or `val` mode, users are able to customize different augmentations for each mode.
|
|
batch (int, optional): Size of batches, this is for `rect`. Defaults to None.
|
|
"""
|
|
return RTDETRDataset(
|
|
img_path=img_path,
|
|
imgsz=self.args.imgsz,
|
|
batch_size=batch,
|
|
augment=mode == 'train', # no augmentation
|
|
hyp=self.args,
|
|
rect=False, # no rect
|
|
cache=self.args.cache or None,
|
|
prefix=colorstr(f'{mode}: '),
|
|
data=self.data)
|
|
|
|
def get_validator(self):
|
|
"""Returns a DetectionValidator for RTDETR model validation."""
|
|
self.loss_names = 'giou_loss', 'cls_loss', 'l1_loss'
|
|
return RTDETRValidator(self.test_loader, save_dir=self.save_dir, args=copy(self.args))
|
|
|
|
def preprocess_batch(self, batch):
|
|
"""Preprocesses a batch of images by scaling and converting to float."""
|
|
batch = super().preprocess_batch(batch)
|
|
bs = len(batch['img'])
|
|
batch_idx = batch['batch_idx']
|
|
gt_bbox, gt_class = [], []
|
|
for i in range(bs):
|
|
gt_bbox.append(batch['bboxes'][batch_idx == i].to(batch_idx.device))
|
|
gt_class.append(batch['cls'][batch_idx == i].to(device=batch_idx.device, dtype=torch.long))
|
|
return batch
|
|
|
|
|
|
def train(cfg=DEFAULT_CFG, use_python=False):
|
|
"""Train and optimize RTDETR model given training data and device."""
|
|
model = 'rtdetr-l.yaml'
|
|
data = cfg.data or 'coco128.yaml' # or yolo.ClassificationDataset("mnist")
|
|
device = cfg.device if cfg.device is not None else ''
|
|
|
|
# NOTE: F.grid_sample which is in rt-detr does not support deterministic=True
|
|
# NOTE: amp training causes nan outputs and end with error while doing bipartite graph matching
|
|
args = dict(model=model,
|
|
data=data,
|
|
device=device,
|
|
imgsz=640,
|
|
exist_ok=True,
|
|
batch=4,
|
|
deterministic=False,
|
|
amp=False)
|
|
trainer = RTDETRTrainer(overrides=args)
|
|
trainer.train()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
train()
|