์ฝ˜ํ…์ธ ๋กœ ๊ฑด๋„ˆ๋›ฐ๊ธฐ

์ฐธ์กฐ ultralytics/models/yolo/detect/train.py

์ฐธ๊ณ 

์ด ํŒŒ์ผ์€ https://github.com/ultralytics/ ultralytics/blob/main/ ultralytics/models/ yolo/detect/train .py์—์„œ ํ™•์ธํ•  ์ˆ˜ ์žˆ์Šต๋‹ˆ๋‹ค. ๋ฌธ์ œ๋ฅผ ๋ฐœ๊ฒฌํ•˜๋ฉด ํ’€ ๋ฆฌํ€˜์ŠคํŠธ (๐Ÿ› ๏ธ)๋ฅผ ์ œ์ถœํ•˜์—ฌ ๋ฌธ์ œ๋ฅผ ํ•ด๊ฒฐํ•˜๋„๋ก ๋„์™€์ฃผ์„ธ์š”. ๊ฐ์‚ฌํ•ฉ๋‹ˆ๋‹ค ๐Ÿ™!



ultralytics.models.yolo.detect.train.DetectionTrainer

๊ธฐ์ง€: BaseTrainer

ํƒ์ง€ ๋ชจ๋ธ์— ๊ธฐ๋ฐ˜ํ•œ ํ›ˆ๋ จ์„ ์œ„ํ•ด BaseTrainer ํด๋ž˜์Šค๋ฅผ ํ™•์žฅํ•œ ํด๋ž˜์Šค์ž…๋‹ˆ๋‹ค.

์˜ˆ์ œ
from ultralytics.models.yolo.detect import DetectionTrainer

args = dict(model='yolov8n.pt', data='coco8.yaml', epochs=3)
trainer = DetectionTrainer(overrides=args)
trainer.train()
์˜ ์†Œ์Šค ์ฝ”๋“œ ultralytics/models/yolo/detect/train.py
class DetectionTrainer(BaseTrainer):
    """
    A class extending the BaseTrainer class for training based on a detection model.

    Example:
        ```python
        from ultralytics.models.yolo.detect import DetectionTrainer

        args = dict(model='yolov8n.pt', data='coco8.yaml', epochs=3)
        trainer = DetectionTrainer(overrides=args)
        trainer.train()
        ```
    """

    def build_dataset(self, img_path, mode="train", batch=None):
        """
        Build YOLO 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.
        """
        gs = max(int(de_parallel(self.model).stride.max() if self.model else 0), 32)
        return build_yolo_dataset(self.args, img_path, batch, self.data, mode=mode, rect=mode == "val", stride=gs)

    def get_dataloader(self, dataset_path, batch_size=16, rank=0, mode="train"):
        """Construct and return dataloader."""
        assert mode in ["train", "val"]
        with torch_distributed_zero_first(rank):  # init dataset *.cache only once if DDP
            dataset = self.build_dataset(dataset_path, mode, batch_size)
        shuffle = mode == "train"
        if getattr(dataset, "rect", False) and shuffle:
            LOGGER.warning("WARNING โš ๏ธ 'rect=True' is incompatible with DataLoader shuffle, setting shuffle=False")
            shuffle = False
        workers = self.args.workers if mode == "train" else self.args.workers * 2
        return build_dataloader(dataset, batch_size, workers, shuffle, rank)  # return dataloader

    def preprocess_batch(self, batch):
        """Preprocesses a batch of images by scaling and converting to float."""
        batch["img"] = batch["img"].to(self.device, non_blocking=True).float() / 255
        if self.args.multi_scale:
            imgs = batch["img"]
            sz = (
                random.randrange(self.args.imgsz * 0.5, self.args.imgsz * 1.5 + self.stride)
                // self.stride
                * self.stride
            )  # size
            sf = sz / max(imgs.shape[2:])  # scale factor
            if sf != 1:
                ns = [
                    math.ceil(x * sf / self.stride) * self.stride for x in imgs.shape[2:]
                ]  # new shape (stretched to gs-multiple)
                imgs = nn.functional.interpolate(imgs, size=ns, mode="bilinear", align_corners=False)
            batch["img"] = imgs
        return batch

    def set_model_attributes(self):
        """Nl = de_parallel(self.model).model[-1].nl  # number of detection layers (to scale hyps)."""
        # self.args.box *= 3 / nl  # scale to layers
        # self.args.cls *= self.data["nc"] / 80 * 3 / nl  # scale to classes and layers
        # self.args.cls *= (self.args.imgsz / 640) ** 2 * 3 / nl  # scale to image size and layers
        self.model.nc = self.data["nc"]  # attach number of classes to model
        self.model.names = self.data["names"]  # attach class names to model
        self.model.args = self.args  # attach hyperparameters to model
        # TODO: self.model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc

    def get_model(self, cfg=None, weights=None, verbose=True):
        """Return a YOLO detection model."""
        model = DetectionModel(cfg, nc=self.data["nc"], verbose=verbose and RANK == -1)
        if weights:
            model.load(weights)
        return model

    def get_validator(self):
        """Returns a DetectionValidator for YOLO model validation."""
        self.loss_names = "box_loss", "cls_loss", "dfl_loss"
        return yolo.detect.DetectionValidator(
            self.test_loader, save_dir=self.save_dir, args=copy(self.args), _callbacks=self.callbacks
        )

    def label_loss_items(self, loss_items=None, prefix="train"):
        """
        Returns a loss dict with labelled training loss items tensor.

        Not needed for classification but necessary for segmentation & detection
        """
        keys = [f"{prefix}/{x}" for x in self.loss_names]
        if loss_items is not None:
            loss_items = [round(float(x), 5) for x in loss_items]  # convert tensors to 5 decimal place floats
            return dict(zip(keys, loss_items))
        else:
            return keys

    def progress_string(self):
        """Returns a formatted string of training progress with epoch, GPU memory, loss, instances and size."""
        return ("\n" + "%11s" * (4 + len(self.loss_names))) % (
            "Epoch",
            "GPU_mem",
            *self.loss_names,
            "Instances",
            "Size",
        )

    def plot_training_samples(self, batch, ni):
        """Plots training samples with their annotations."""
        plot_images(
            images=batch["img"],
            batch_idx=batch["batch_idx"],
            cls=batch["cls"].squeeze(-1),
            bboxes=batch["bboxes"],
            paths=batch["im_file"],
            fname=self.save_dir / f"train_batch{ni}.jpg",
            on_plot=self.on_plot,
        )

    def plot_metrics(self):
        """Plots metrics from a CSV file."""
        plot_results(file=self.csv, on_plot=self.on_plot)  # save results.png

    def plot_training_labels(self):
        """Create a labeled training plot of the YOLO model."""
        boxes = np.concatenate([lb["bboxes"] for lb in self.train_loader.dataset.labels], 0)
        cls = np.concatenate([lb["cls"] for lb in self.train_loader.dataset.labels], 0)
        plot_labels(boxes, cls.squeeze(), names=self.data["names"], save_dir=self.save_dir, on_plot=self.on_plot)

build_dataset(img_path, mode='train', batch=None)

YOLO ๋ฐ์ดํ„ฐ ์ง‘ํ•ฉ์„ ๊ตฌ์ถ•ํ•ฉ๋‹ˆ๋‹ค.

๋งค๊ฐœ๋ณ€์ˆ˜:

์ด๋ฆ„ ์œ ํ˜• ์„ค๋ช… ๊ธฐ๋ณธ๊ฐ’
img_path str

์ด๋ฏธ์ง€๊ฐ€ ํฌํ•จ๋œ ํด๋”์˜ ๊ฒฝ๋กœ์ž…๋‹ˆ๋‹ค.

ํ•„์ˆ˜
mode str

train ๋ชจ๋“œ ๋˜๋Š” val ๋ชจ๋“œ์—์„œ ์‚ฌ์šฉ์ž๋Š” ๊ฐ ๋ชจ๋“œ์— ๋Œ€ํ•ด ์„œ๋กœ ๋‹ค๋ฅธ ์ฆ๊ฐ• ๊ธฐ๋Šฅ์„ ์‚ฌ์šฉ์ž ์ง€์ •ํ•  ์ˆ˜ ์žˆ์Šต๋‹ˆ๋‹ค.

'train'
batch int

๋ฐฐ์น˜์˜ ํฌ๊ธฐ, ์ด๊ฒƒ์€ ๋‹ค์Œ๊ณผ ๊ฐ™์Šต๋‹ˆ๋‹ค. rect. ๊ธฐ๋ณธ๊ฐ’์€ ์—†์Œ์ž…๋‹ˆ๋‹ค.

None
์˜ ์†Œ์Šค ์ฝ”๋“œ ultralytics/models/yolo/detect/train.py
def build_dataset(self, img_path, mode="train", batch=None):
    """
    Build YOLO 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.
    """
    gs = max(int(de_parallel(self.model).stride.max() if self.model else 0), 32)
    return build_yolo_dataset(self.args, img_path, batch, self.data, mode=mode, rect=mode == "val", stride=gs)

get_dataloader(dataset_path, batch_size=16, rank=0, mode='train')

๋ฐ์ดํ„ฐ๋กœ๋”๋ฅผ ์ƒ์„ฑํ•˜๊ณ  ๋ฐ˜ํ™˜ํ•ฉ๋‹ˆ๋‹ค.

์˜ ์†Œ์Šค ์ฝ”๋“œ ultralytics/models/yolo/detect/train.py
def get_dataloader(self, dataset_path, batch_size=16, rank=0, mode="train"):
    """Construct and return dataloader."""
    assert mode in ["train", "val"]
    with torch_distributed_zero_first(rank):  # init dataset *.cache only once if DDP
        dataset = self.build_dataset(dataset_path, mode, batch_size)
    shuffle = mode == "train"
    if getattr(dataset, "rect", False) and shuffle:
        LOGGER.warning("WARNING โš ๏ธ 'rect=True' is incompatible with DataLoader shuffle, setting shuffle=False")
        shuffle = False
    workers = self.args.workers if mode == "train" else self.args.workers * 2
    return build_dataloader(dataset, batch_size, workers, shuffle, rank)  # return dataloader

get_model(cfg=None, weights=None, verbose=True)

YOLO ํƒ์ง€ ๋ชจ๋ธ์„ ๋ฐ˜ํ™˜ํ•ฉ๋‹ˆ๋‹ค.

์˜ ์†Œ์Šค ์ฝ”๋“œ ultralytics/models/yolo/detect/train.py
def get_model(self, cfg=None, weights=None, verbose=True):
    """Return a YOLO detection model."""
    model = DetectionModel(cfg, nc=self.data["nc"], verbose=verbose and RANK == -1)
    if weights:
        model.load(weights)
    return model

get_validator()

YOLO ๋ชจ๋ธ ์œ ํšจ์„ฑ ๊ฒ€์‚ฌ๋ฅผ ์œ„ํ•œ DetectionValidator๋ฅผ ๋ฐ˜ํ™˜ํ•ฉ๋‹ˆ๋‹ค.

์˜ ์†Œ์Šค ์ฝ”๋“œ ultralytics/models/yolo/detect/train.py
def get_validator(self):
    """Returns a DetectionValidator for YOLO model validation."""
    self.loss_names = "box_loss", "cls_loss", "dfl_loss"
    return yolo.detect.DetectionValidator(
        self.test_loader, save_dir=self.save_dir, args=copy(self.args), _callbacks=self.callbacks
    )

label_loss_items(loss_items=None, prefix='train')

๋ ˆ์ด๋ธ”์ด ์ง€์ •๋œ ํ›ˆ๋ จ ์†์‹ค ํ•ญ๋ชฉ์ด ์žˆ๋Š” ์†์‹ค ๋”•์…”๋„ˆ๋ฆฌ๋ฅผ ๋ฐ˜ํ™˜ํ•ฉ๋‹ˆ๋‹ค tensor.

๋ถ„๋ฅ˜์—๋Š” ํ•„์š”ํ•˜์ง€ ์•Š์ง€๋งŒ ์„ธ๋ถ„ํ™” ๋ฐ ํƒ์ง€์—๋Š” ํ•„์š”ํ•ฉ๋‹ˆ๋‹ค.

์˜ ์†Œ์Šค ์ฝ”๋“œ ultralytics/models/yolo/detect/train.py
def label_loss_items(self, loss_items=None, prefix="train"):
    """
    Returns a loss dict with labelled training loss items tensor.

    Not needed for classification but necessary for segmentation & detection
    """
    keys = [f"{prefix}/{x}" for x in self.loss_names]
    if loss_items is not None:
        loss_items = [round(float(x), 5) for x in loss_items]  # convert tensors to 5 decimal place floats
        return dict(zip(keys, loss_items))
    else:
        return keys

plot_metrics()

CSV ํŒŒ์ผ์—์„œ ๋ฉ”ํŠธ๋ฆญ์„ ํ”Œ๋กœํŒ…ํ•ฉ๋‹ˆ๋‹ค.

์˜ ์†Œ์Šค ์ฝ”๋“œ ultralytics/models/yolo/detect/train.py
def plot_metrics(self):
    """Plots metrics from a CSV file."""
    plot_results(file=self.csv, on_plot=self.on_plot)  # save results.png

plot_training_labels()

YOLO ๋ชจ๋ธ์˜ ๋ ˆ์ด๋ธ”์ด ์ง€์ •๋œ ํŠธ๋ ˆ์ด๋‹ ํ”Œ๋กฏ์„ ๋งŒ๋“ญ๋‹ˆ๋‹ค.

์˜ ์†Œ์Šค ์ฝ”๋“œ ultralytics/models/yolo/detect/train.py
def plot_training_labels(self):
    """Create a labeled training plot of the YOLO model."""
    boxes = np.concatenate([lb["bboxes"] for lb in self.train_loader.dataset.labels], 0)
    cls = np.concatenate([lb["cls"] for lb in self.train_loader.dataset.labels], 0)
    plot_labels(boxes, cls.squeeze(), names=self.data["names"], save_dir=self.save_dir, on_plot=self.on_plot)

plot_training_samples(batch, ni)

ํŠธ๋ ˆ์ด๋‹ ์ƒ˜ํ”Œ์„ ์ฃผ์„๊ณผ ํ•จ๊ป˜ ํ”Œ๋กฏํ•ฉ๋‹ˆ๋‹ค.

์˜ ์†Œ์Šค ์ฝ”๋“œ ultralytics/models/yolo/detect/train.py
def plot_training_samples(self, batch, ni):
    """Plots training samples with their annotations."""
    plot_images(
        images=batch["img"],
        batch_idx=batch["batch_idx"],
        cls=batch["cls"].squeeze(-1),
        bboxes=batch["bboxes"],
        paths=batch["im_file"],
        fname=self.save_dir / f"train_batch{ni}.jpg",
        on_plot=self.on_plot,
    )

preprocess_batch(batch)

์ด๋ฏธ์ง€์˜ ํฌ๊ธฐ๋ฅผ ์กฐ์ •ํ•˜๊ณ  ํ”Œ๋กœํŠธ๋กœ ๋ณ€ํ™˜ํ•˜์—ฌ ์ด๋ฏธ์ง€ ๋ฐฐ์น˜๋ฅผ ์ „์ฒ˜๋ฆฌํ•ฉ๋‹ˆ๋‹ค.

์˜ ์†Œ์Šค ์ฝ”๋“œ ultralytics/models/yolo/detect/train.py
def preprocess_batch(self, batch):
    """Preprocesses a batch of images by scaling and converting to float."""
    batch["img"] = batch["img"].to(self.device, non_blocking=True).float() / 255
    if self.args.multi_scale:
        imgs = batch["img"]
        sz = (
            random.randrange(self.args.imgsz * 0.5, self.args.imgsz * 1.5 + self.stride)
            // self.stride
            * self.stride
        )  # size
        sf = sz / max(imgs.shape[2:])  # scale factor
        if sf != 1:
            ns = [
                math.ceil(x * sf / self.stride) * self.stride for x in imgs.shape[2:]
            ]  # new shape (stretched to gs-multiple)
            imgs = nn.functional.interpolate(imgs, size=ns, mode="bilinear", align_corners=False)
        batch["img"] = imgs
    return batch

progress_string()

์—ํฌํฌ, GPU ๋ฉ”๋ชจ๋ฆฌ, ์†์‹ค, ์ธ์Šคํ„ด์Šค ๋ฐ ํฌ๊ธฐ๊ฐ€ ํฌํ•จ๋œ ํ›ˆ๋ จ ์ง„ํ–‰๋ฅ ์˜ ํ˜•์‹ํ™”๋œ ๋ฌธ์ž์—ด์„ ๋ฐ˜ํ™˜ํ•ฉ๋‹ˆ๋‹ค.

์˜ ์†Œ์Šค ์ฝ”๋“œ ultralytics/models/yolo/detect/train.py
def progress_string(self):
    """Returns a formatted string of training progress with epoch, GPU memory, loss, instances and size."""
    return ("\n" + "%11s" * (4 + len(self.loss_names))) % (
        "Epoch",
        "GPU_mem",
        *self.loss_names,
        "Instances",
        "Size",
    )

set_model_attributes()

Nl = de_parallel(self.model).model[-1].nl # ํƒ์ง€ ๋ ˆ์ด์–ด ์ˆ˜(๊ฐ€์„ค์„ ํ™•์žฅํ•˜๊ธฐ ์œ„ํ•œ).

์˜ ์†Œ์Šค ์ฝ”๋“œ ultralytics/models/yolo/detect/train.py
def set_model_attributes(self):
    """Nl = de_parallel(self.model).model[-1].nl  # number of detection layers (to scale hyps)."""
    # self.args.box *= 3 / nl  # scale to layers
    # self.args.cls *= self.data["nc"] / 80 * 3 / nl  # scale to classes and layers
    # self.args.cls *= (self.args.imgsz / 640) ** 2 * 3 / nl  # scale to image size and layers
    self.model.nc = self.data["nc"]  # attach number of classes to model
    self.model.names = self.data["names"]  # attach class names to model
    self.model.args = self.args  # attach hyperparameters to model





์ƒ์„ฑ๋จ 2023-11-12, ์—…๋ฐ์ดํŠธ๋จ 2023-11-25
์ž‘์„ฑ์ž: glenn-jocher (3)