Skip to content

Commit

Permalink
update code
Browse files Browse the repository at this point in the history
  • Loading branch information
wondervictor committed Feb 1, 2024
1 parent 1325194 commit dc17c77
Show file tree
Hide file tree
Showing 43 changed files with 5,860 additions and 7 deletions.
1,347 changes: 1,347 additions & 0 deletions LICENSE

Large diffs are not rendered by default.

23 changes: 16 additions & 7 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -27,11 +27,20 @@

## Updates

YOLO-World is under active development and please stay tuned ☕️!

`[2024-1-31]:` We are excited to launch **YOLO-World**, a cutting-edge real-time open-vocabulary object detector.
`[2024-2-1]:` We deploy the YOLO-World demo on [HuggingFace 🤗](https://huggingface.co/spaces/stevengrove/YOLO-World), you can try it now!
`[2024-2-1]:` We deploy the YOLO-World demo on [HuggingFace 🤗](https://huggingface.co/spaces/stevengrove/YOLO-World), you can try it now!
`[2024-2-1]:` We've released the code and weights of YOLO-World now!

## TODO

- [ ] Complete documents for pre-training YOLO-World.
- [ ] Deployment toolkits, e.g., ONNX or TensorRT.
- [ ] Inference acceleration and scripts for speed evaluation.
- [ ] Automatic labeling framework for image-text pairs, such as CC3M.


## Highlights

This repo contains the PyTorch implementation, pre-trained weights, and pre-training/fine-tuning code for YOLO-World.
Expand Down Expand Up @@ -62,14 +71,14 @@ We've pre-trained YOLO-World-S/M/L from scratch and evaluate on the `LVIS val-1.

### Zero-shot Inference on LVIS dataset

| model | Pre-train Data | AP | AP<sub>r</sub> | AP<sub>c</sub> | AP<sub>f</sub> | weights | log |
| :---- | :------------- | :-:| :------------: |:-------------: | :-------: | :---: | :---: |
| [YOLO-World-S]() | O365+GoldG | 26.2 | 19.1 | 23.6 | 29.8 | - | [coming soon] | [log]() |
| [YOLO-World-M]() | O365+GoldG | 31.0 | 23.8 | 29.2 | 33.9 | - | [coming soon] | [log]() |
| [YOLO-World-L]() | O365+GoldG | 35.0 | 27.1 | 32.8 | 38.3 | - | [coming soon]| [log]() |
| model | Pre-train Data | AP<sup>mini</su> | AP<sub>r</sub> | AP<sub>c</sub> | AP<sub>f</sub> | AP<sup>val</su> | AP<sub>r</sub> | AP<sub>c</sub> | AP<sub>f</sub> | weights | log |
| :---- | :------------- | :-:| :------------: |:-------------: | :------------: | :-:| :------------: |:-------------: | :------------: | :---: | :---: |
| [YOLO-World-S]() | O365+GoldG | 26.2 | 19.1 | 23.6 | 29.8 | - | | [log]() |
| [YOLO-World-M]() | O365+GoldG | 31.0 | 23.8 | 29.2 | 33.9 | - | | [log]() |
| [YOLO-World-L]() | O365+GoldG | 35.0 | 27.1 | 32.8 | 38.3 | - | | [log]() |

**NOTE:**
1. The evaluation results are tested on LVIS minival in a zero-shot manner.
1. The evaluation results are tested on LVIS minival ([fixed AP](https://github.com/achalddave/large-vocab-devil) with 1,000 dets) and LVIS val in a zero-shot manner.


### Finetuning on COCO dataset
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,173 @@
_base_ = ('../../third_party/mmyolo/configs/yolov8/'
'yolov8_l_syncbn_fast_8xb16-500e_coco.py')
custom_imports = dict(imports=['yolo_world'],
allow_failed_imports=False)

# hyper-parameters
num_classes = 1203
num_training_classes = 80
max_epochs = 100 # Maximum training epochs
close_mosaic_epochs = 2
save_epoch_intervals = 2
text_channels = 512
neck_embed_channels = [128, 256, _base_.last_stage_out_channels // 2]
neck_num_heads = [4, 8, _base_.last_stage_out_channels // 2 // 32]
base_lr = 2e-3
weight_decay = 0.05 / 2
train_batch_size_per_gpu = 16

# model settings
model = dict(
type='YOLOWorldDetector',
mm_neck=True,
num_train_classes=num_training_classes,
num_test_classes=num_classes,
data_preprocessor=dict(type='YOLOWDetDataPreprocessor'),
backbone=dict(
_delete_=True,
type='MultiModalYOLOBackbone',
image_model={{_base_.model.backbone}},
text_model=dict(
type='HuggingCLIPLanguageBackbone',
model_name='pretrained_models/clip-vit-base-patch32-projection',
frozen_modules=['all'])),
neck=dict(type='YOLOWolrdDualPAFPN',
guide_channels=text_channels,
embed_channels=neck_embed_channels,
num_heads=neck_num_heads,
block_cfg=dict(type='MaxSigmoidCSPLayerWithTwoConv'),
text_enhancder=dict(type='ImagePoolingAttentionModule',
embed_channels=256,
num_heads=8)),
bbox_head=dict(type='YOLOWorldHead',
head_module=dict(type='YOLOWorldHeadModule',
embed_dims=text_channels,
num_classes=num_training_classes)),
train_cfg=dict(assigner=dict(num_classes=num_training_classes)))

# dataset settings
text_transform = [
dict(type='RandomLoadText',
num_neg_samples=(num_classes, num_classes),
max_num_samples=num_training_classes,
padding_to_max=True,
padding_value=''),
dict(type='mmdet.PackDetInputs',
meta_keys=('img_id', 'img_path', 'ori_shape', 'img_shape', 'flip',
'flip_direction', 'texts'))
]
train_pipeline = [
*_base_.pre_transform,
dict(type='MultiModalMosaic',
img_scale=_base_.img_scale,
pad_val=114.0,
pre_transform=_base_.pre_transform),
dict(
type='YOLOv5RandomAffine',
max_rotate_degree=0.0,
max_shear_degree=0.0,
scaling_ratio_range=(1 - _base_.affine_scale, 1 + _base_.affine_scale),
max_aspect_ratio=_base_.max_aspect_ratio,
border=(-_base_.img_scale[0] // 2, -_base_.img_scale[1] // 2),
border_val=(114, 114, 114)),
*_base_.last_transform[:-1],
*text_transform,
]
train_pipeline_stage2 = [*_base_.train_pipeline_stage2[:-1], *text_transform]
obj365v1_train_dataset = dict(
type='MultiModalDataset',
dataset=dict(
type='YOLOv5Objects365V1Dataset',
data_root='data/objects365v1/',
ann_file='annotations/objects365_train.json',
data_prefix=dict(img='train/'),
filter_cfg=dict(filter_empty_gt=False, min_size=32)),
class_text_path='data/captions/obj365v1_class_captions.json',
pipeline=train_pipeline)

mg_train_dataset = dict(type='YOLOv5MixedGroundingDataset',
data_root='data/mixed_grounding/',
ann_file='annotations/final_mixed_train_no_coco.json',
data_prefix=dict(img='gqa/images/'),
filter_cfg=dict(filter_empty_gt=False, min_size=32),
pipeline=train_pipeline)

flickr_train_dataset = dict(
type='YOLOv5MixedGroundingDataset',
data_root='data/flickr/',
ann_file='annotations/final_flickr_separateGT_train.json',
data_prefix=dict(img='full_images/'),
filter_cfg=dict(filter_empty_gt=True, min_size=32),
pipeline=train_pipeline)

train_dataloader = dict(batch_size=train_batch_size_per_gpu,
collate_fn=dict(type='yolow_collate'),
dataset=dict(_delete_=True,
type='ConcatDataset',
datasets=[
obj365v1_train_dataset,
flickr_train_dataset, mg_train_dataset
],
ignore_keys=['classes', 'palette']))

test_pipeline = [
*_base_.test_pipeline[:-1],
dict(type='LoadText'),
dict(type='mmdet.PackDetInputs',
meta_keys=('img_id', 'img_path', 'ori_shape', 'img_shape',
'scale_factor', 'pad_param', 'texts'))
]
coco_val_dataset = dict(
_delete_=True,
type='MultiModalDataset',
dataset=dict(type='YOLOv5LVISV1Dataset',
data_root='data/coco/',
test_mode=True,
ann_file='lvis/lvis_v1_minival_inserted_image_name.json',
data_prefix=dict(img=''),
batch_shapes_cfg=None),
class_text_path='data/captions/lvis_v1_class_captions.json',
pipeline=test_pipeline)
val_dataloader = dict(dataset=coco_val_dataset)
test_dataloader = val_dataloader

val_evaluator = dict(type='mmdet.LVISMetric',
ann_file='data/coco/lvis/\
lvis_v1_minival_inserted_image_name.json',
metric='bbox')
test_evaluator = val_evaluator

# training settings
default_hooks = dict(param_scheduler=dict(max_epochs=max_epochs),
checkpoint=dict(interval=save_epoch_intervals,
rule='greater'))
custom_hooks = [
dict(type='EMAHook',
ema_type='ExpMomentumEMA',
momentum=0.0001,
update_buffers=True,
strict_load=False,
priority=49),
dict(type='mmdet.PipelineSwitchHook',
switch_epoch=max_epochs - close_mosaic_epochs,
switch_pipeline=train_pipeline_stage2)
]
train_cfg = dict(max_epochs=max_epochs,
val_interval=10,
dynamic_intervals=[((max_epochs - close_mosaic_epochs),
_base_.val_interval_stage2)])
optim_wrapper = dict(optimizer=dict(
_delete_=True,
type='AdamW',
lr=base_lr,
weight_decay=weight_decay,
batch_size_per_gpu=train_batch_size_per_gpu),
paramwise_cfg=dict(bias_decay_mult=0.0,
norm_decay_mult=0.0,
custom_keys={
'backbone.text_model':
dict(lr_mult=0.01),
'logit_scale':
dict(weight_decay=0.0)
}),
constructor='YOLOWv5OptimizerConstructor')
Original file line number Diff line number Diff line change
@@ -0,0 +1,173 @@
_base_ = ('../../third_party/mmyolo/configs/yolov8/'
'yolov8_m_syncbn_fast_8xb16-500e_coco.py')
custom_imports = dict(imports=['yolo_world'],
allow_failed_imports=False)

# hyper-parameters
num_classes = 1203
num_training_classes = 80
max_epochs = 100 # Maximum training epochs
close_mosaic_epochs = 2
save_epoch_intervals = 2
text_channels = 512
neck_embed_channels = [128, 256, _base_.last_stage_out_channels // 2]
neck_num_heads = [4, 8, _base_.last_stage_out_channels // 2 // 32]
base_lr = 2e-3
weight_decay = 0.05 / 2
train_batch_size_per_gpu = 16

# model settings
model = dict(
type='YOLOWorldDetector',
mm_neck=True,
num_train_classes=num_training_classes,
num_test_classes=num_classes,
data_preprocessor=dict(type='YOLOWDetDataPreprocessor'),
backbone=dict(
_delete_=True,
type='MultiModalYOLOBackbone',
image_model={{_base_.model.backbone}},
text_model=dict(
type='HuggingCLIPLanguageBackbone',
model_name='pretrained_models/clip-vit-base-patch32-projection',
frozen_modules=['all'])),
neck=dict(type='YOLOWolrdDualPAFPN',
guide_channels=text_channels,
embed_channels=neck_embed_channels,
num_heads=neck_num_heads,
block_cfg=dict(type='MaxSigmoidCSPLayerWithTwoConv'),
text_enhancder=dict(type='ImagePoolingAttentionModule',
embed_channels=256,
num_heads=8)),
bbox_head=dict(type='YOLOWorldHead',
head_module=dict(type='YOLOWorldHeadModule',
embed_dims=text_channels,
num_classes=num_training_classes)),
train_cfg=dict(assigner=dict(num_classes=num_training_classes)))

# dataset settings
text_transform = [
dict(type='RandomLoadText',
num_neg_samples=(num_classes, num_classes),
max_num_samples=num_training_classes,
padding_to_max=True,
padding_value=''),
dict(type='mmdet.PackDetInputs',
meta_keys=('img_id', 'img_path', 'ori_shape', 'img_shape', 'flip',
'flip_direction', 'texts'))
]
train_pipeline = [
*_base_.pre_transform,
dict(type='MultiModalMosaic',
img_scale=_base_.img_scale,
pad_val=114.0,
pre_transform=_base_.pre_transform),
dict(
type='YOLOv5RandomAffine',
max_rotate_degree=0.0,
max_shear_degree=0.0,
scaling_ratio_range=(1 - _base_.affine_scale, 1 + _base_.affine_scale),
max_aspect_ratio=_base_.max_aspect_ratio,
border=(-_base_.img_scale[0] // 2, -_base_.img_scale[1] // 2),
border_val=(114, 114, 114)),
*_base_.last_transform[:-1],
*text_transform,
]
train_pipeline_stage2 = [*_base_.train_pipeline_stage2[:-1], *text_transform]
obj365v1_train_dataset = dict(
type='MultiModalDataset',
dataset=dict(
type='YOLOv5Objects365V1Dataset',
data_root='data/objects365v1/',
ann_file='annotations/objects365_train.json',
data_prefix=dict(img='train/'),
filter_cfg=dict(filter_empty_gt=False, min_size=32)),
class_text_path='data/captions/obj365v1_class_captions.json',
pipeline=train_pipeline)

mg_train_dataset = dict(type='YOLOv5MixedGroundingDataset',
data_root='data/mixed_grounding/',
ann_file='annotations/final_mixed_train_no_coco.json',
data_prefix=dict(img='gqa/images/'),
filter_cfg=dict(filter_empty_gt=False, min_size=32),
pipeline=train_pipeline)

flickr_train_dataset = dict(
type='YOLOv5MixedGroundingDataset',
data_root='data/flickr/',
ann_file='annotations/final_flickr_separateGT_train.json',
data_prefix=dict(img='full_images/'),
filter_cfg=dict(filter_empty_gt=True, min_size=32),
pipeline=train_pipeline)

train_dataloader = dict(batch_size=train_batch_size_per_gpu,
collate_fn=dict(type='yolow_collate'),
dataset=dict(_delete_=True,
type='ConcatDataset',
datasets=[
obj365v1_train_dataset,
flickr_train_dataset, mg_train_dataset
],
ignore_keys=['classes', 'palette']))

test_pipeline = [
*_base_.test_pipeline[:-1],
dict(type='LoadText'),
dict(type='mmdet.PackDetInputs',
meta_keys=('img_id', 'img_path', 'ori_shape', 'img_shape',
'scale_factor', 'pad_param', 'texts'))
]
coco_val_dataset = dict(
_delete_=True,
type='MultiModalDataset',
dataset=dict(type='YOLOv5LVISV1Dataset',
data_root='data/coco/',
test_mode=True,
ann_file='lvis/lvis_v1_minival_inserted_image_name.json',
data_prefix=dict(img=''),
batch_shapes_cfg=None),
class_text_path='data/captions/lvis_v1_class_captions.json',
pipeline=test_pipeline)
val_dataloader = dict(dataset=coco_val_dataset)
test_dataloader = val_dataloader

val_evaluator = dict(type='mmdet.LVISMetric',
ann_file='data/coco/lvis/\
lvis_v1_minival_inserted_image_name.json',
metric='bbox')
test_evaluator = val_evaluator

# training settings
default_hooks = dict(param_scheduler=dict(max_epochs=max_epochs),
checkpoint=dict(interval=save_epoch_intervals,
rule='greater'))
custom_hooks = [
dict(type='EMAHook',
ema_type='ExpMomentumEMA',
momentum=0.0001,
update_buffers=True,
strict_load=False,
priority=49),
dict(type='mmdet.PipelineSwitchHook',
switch_epoch=max_epochs - close_mosaic_epochs,
switch_pipeline=train_pipeline_stage2)
]
train_cfg = dict(max_epochs=max_epochs,
val_interval=10,
dynamic_intervals=[((max_epochs - close_mosaic_epochs),
_base_.val_interval_stage2)])
optim_wrapper = dict(optimizer=dict(
_delete_=True,
type='AdamW',
lr=base_lr,
weight_decay=weight_decay,
batch_size_per_gpu=train_batch_size_per_gpu),
paramwise_cfg=dict(bias_decay_mult=0.0,
norm_decay_mult=0.0,
custom_keys={
'backbone.text_model':
dict(lr_mult=0.01),
'logit_scale':
dict(weight_decay=0.0)
}),
constructor='YOLOWv5OptimizerConstructor')
Loading

0 comments on commit dc17c77

Please sign in to comment.