Skip to content

Reference for ultralytics/models/yolo/world/train.py

Note

This file is available at https://github.com/ultralytics/ultralytics/blob/main/ultralytics/models/yolo/world/train.py. If you spot a problem please help fix it by contributing a Pull Request 🛠️. Thank you 🙏!


ultralytics.models.yolo.world.train.WorldTrainer

WorldTrainer(cfg=DEFAULT_CFG, overrides=None, _callbacks=None)

Bases: DetectionTrainer

A class to fine-tune a world model on a close-set dataset.

Example 
from ultralytics.models.yolo.world import WorldModel

args = dict(model="yolov8s-world.pt", data="coco8.yaml", epochs=3)
trainer = WorldTrainer(overrides=args)
trainer.train()
Source code in ultralytics/models/yolo/world/train.py 
38
39
40
41
42
43
44
45
46
47
48
49
50
def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
    """Initialize a WorldTrainer object with given arguments."""
    if overrides is None:
        overrides = {}
    super().__init__(cfg, overrides, _callbacks)

    # Import and assign clip
    try:
        import clip
    except ImportError:
        checks.check_requirements("git+https://github.com/ultralytics/CLIP.git")
        import clip
    self.clip = clip

build_dataset

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

Build YOLO Dataset.

Parameters:

NameTypeDescriptionDefault
img_pathstr

Path to the folder containing images.

required
modestr

train mode or val mode, users are able to customize different augmentations for each mode.

'train'
batchint

Size of batches, this is for rect. Defaults to None.

None
Source code in ultralytics/models/yolo/world/train.py 
68
69
70
71
72
73
74
75
76
77
78
79
80
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, multi_modal=mode == "train"
    )

get_model

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

Return WorldModel initialized with specified config and weights.

Source code in ultralytics/models/yolo/world/train.py 
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
def get_model(self, cfg=None, weights=None, verbose=True):
    """Return WorldModel initialized with specified config and weights."""
    # NOTE: This `nc` here is the max number of different text samples in one image, rather than the actual `nc`.
    # NOTE: Following the official config, nc hard-coded to 80 for now.
    model = WorldModel(
        cfg["yaml_file"] if isinstance(cfg, dict) else cfg,
        ch=3,
        nc=min(self.data["nc"], 80),
        verbose=verbose and RANK == -1,
    )
    if weights:
        model.load(weights)
    self.add_callback("on_pretrain_routine_end", on_pretrain_routine_end)

    return model

preprocess_batch

preprocess_batch(batch)

Preprocesses a batch of images for YOLOWorld training, adjusting formatting and dimensions as needed.

Source code in ultralytics/models/yolo/world/train.py 
82
83
84
85
86
87
88
89
90
91
92
def preprocess_batch(self, batch):
    """Preprocesses a batch of images for YOLOWorld training, adjusting formatting and dimensions as needed."""
    batch = super().preprocess_batch(batch)

    # NOTE: add text features
    texts = list(itertools.chain(*batch["texts"]))
    text_token = self.clip.tokenize(texts).to(batch["img"].device)
    txt_feats = self.text_model.encode_text(text_token).to(dtype=batch["img"].dtype)  # torch.float32
    txt_feats = txt_feats / txt_feats.norm(p=2, dim=-1, keepdim=True)
    batch["txt_feats"] = txt_feats.reshape(len(batch["texts"]), -1, txt_feats.shape[-1])
    return batch




ultralytics.models.yolo.world.train.on_pretrain_routine_end

on_pretrain_routine_end(trainer)

Callback.

Source code in ultralytics/models/yolo/world/train.py 
12
13
14
15
16
17
18
19
20
21
def on_pretrain_routine_end(trainer):
    """Callback."""
    if RANK in {-1, 0}:
        # NOTE: for evaluation
        names = [name.split("/")[0] for name in list(trainer.test_loader.dataset.data["names"].values())]
        de_parallel(trainer.ema.ema).set_classes(names, cache_clip_model=False)
    device = next(trainer.model.parameters()).device
    trainer.text_model, _ = trainer.clip.load("ViT-B/32", device=device)
    for p in trainer.text_model.parameters():
        p.requires_grad_(False)



📅 Created 7 months ago ✏️ Updated 1 month ago
jk4e glenn-jocher Burhan-Q Laughing-q