Segmentation COCO minimum example

[23pointsNorth]

💡 Your Question

I am trying to recreate a minimal example for training a segmentation network (is YOLO-NAS going to support this in the future?) with COCO dataset.

I am successfully loading the dataset, but having issues during training. Relevant code based on the supervisely segmentation quickstart notebook:

from super_gradients import Trainer, setup_device
from super_gradients.training import  MultiGPUMode, dataloaders
import os

CHECKPOINT_DIR = 'ckpts/'
os.makedirs(CHECKPOINT_DIR, exist_ok=True)
setup_device(device="cuda", multi_gpu=MultiGPUMode.OFF)
trainer = Trainer(experiment_name="segmentation_quick_start", ckpt_root_dir=CHECKPOINT_DIR)

# Data
batch_size = 8
train_loader = dataloaders.coco_segmentation_train(dataset_params={"root_dir": "data/coco"}, dataloader_params={"batch_size": batch_size})
valid_loader = dataloaders.coco_segmentation_val(dataset_params={"root_dir": "data/coco"}, dataloader_params={"batch_size": batch_size})
from prettyformatter import pprint

print('Dataloader parameters:')
pprint(train_loader.dataloader_params)
print('Dataset parameters')
pprint(train_loader.dataset.dataset_params)

# Training
from super_gradients.training import models
from super_gradients.common.object_names import Models

model = models.get(model_name=Models.PP_LITE_T_SEG,
                   arch_params={"use_aux_heads": False},
                   num_classes=len(train_loader.dataset.dataset_classes_inclusion_tuples_list)) # or -1, because it has background?
from super_gradients.training.metrics.segmentation_metrics import BinaryIOU
from super_gradients.training.utils.callbacks import BinarySegmentationVisualizationCallback, Phase

train_params = {"max_epochs": 30,
                "lr_mode": "cosine",
                "initial_lr": 0.01,
                "lr_warmup_epochs": 5,
                "multiply_head_lr": 10,
                "optimizer": "SGD",
                "loss": "bce_dice_loss",
                "ema": True,
                "zero_weight_decay_on_bias_and_bn": True,
                "average_best_models": True,
                "metric_to_watch": "target_IOU",
                "greater_metric_to_watch_is_better": True,
                "train_metrics_list": [BinaryIOU()],
                "valid_metrics_list": [BinaryIOU()],
                "loss_logging_items_names": ["loss"],
                "phase_callbacks": [BinarySegmentationVisualizationCallback(phase=Phase.VALIDATION_BATCH_END,
                                                                            freq=1,
                                                                            last_img_idx_in_batch=4)],
                }
trainer.train(model=model, training_params=train_params, train_loader=train_loader, valid_loader=valid_loader) # << Error happens here

print("Best Checkpoint mIoU is: "+ str(trainer.best_metric.item()))
print(trainer.checkpoints_dir_path)

Error message around different side between model output and dataset target

  File "/home/bot/code/sg/sg_quickstart_segmentation.py", line 206, in <module>
    trainer.train(model=model, training_params=train_params, train_loader=train_loader, valid_loader=valid_loader)
  File "/home/bot/.local/lib/python3.10/site-packages/super_gradients/training/sg_trainer/sg_trainer.py", line 1362, in train
    train_metrics_tuple = self._train_epoch(context=context, silent_mode=silent_mode)
  File "/home/bot/.local/lib/python3.10/site-packages/super_gradients/training/sg_trainer/sg_trainer.py", line 442, in _train_epoch
    loss, loss_log_items = self._get_losses(outputs, targets)
  File "/home/bot/.local/lib/python3.10/site-packages/super_gradients/training/sg_trainer/sg_trainer.py", line 475, in _get_losses
    loss = self.criterion(outputs, targets)
  File "/home/bot/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/bot/.local/lib/python3.10/site-packages/super_gradients/training/losses/bce_dice_loss.py", line 34, in forward
    return self.loss_weights[0] * self.bce(input, target) + self.loss_weights[1] * self.dice(input, target)
  File "/home/bot/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/bot/.local/lib/python3.10/site-packages/super_gradients/training/losses/bce_loss.py", line 16, in forward
    return super(BCE, self).forward(input.squeeze(1), target.float())
  File "/home/bot/.local/lib/python3.10/site-packages/torch/nn/modules/loss.py", line 720, in forward
    return F.binary_cross_entropy_with_logits(input, target,
  File "/home/bot/.local/lib/python3.10/site-packages/torch/nn/functional.py", line 3163, in binary_cross_entropy_with_logits
    raise ValueError("Target size ({}) must be the same as input size ({})".format(target.size(), input.size()))
ValueError: Target size (torch.Size([8, 512, 512])) must be the same as input size (torch.Size([8, 21, 512, 512]))

Would love some help resoling this, plus makes for a great MVP/example.

Versions

Collecting environment information...
PyTorch version: 2.0.1+cu117
Is debug build: False
CUDA used to build PyTorch: 11.7
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.3 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.35

Python version: 3.10.12 (main, Jun 11 2023, 05:26:28) [GCC 11.4.0] (64-bit runtime)
Python platform: Linux-6.2.0-26-generic-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 3050 Ti Laptop GPU
Nvidia driver version: 535.86.05
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                    x86_64
CPU op-mode(s):                  32-bit, 64-bit
Address sizes:                   39 bits physical, 48 bits virtual
Byte Order:                      Little Endian
CPU(s):                          16
On-line CPU(s) list:             0-15
Vendor ID:                       GenuineIntel
Model name:                      11th Gen Intel(R) Core(TM) i7-11800H @ 2.30GHz
CPU family:                      6
Model:                           141
Thread(s) per core:              2
Core(s) per socket:              8
Socket(s):                       1
Stepping:                        1
CPU max MHz:                     4600.0000
CPU min MHz:                     800.0000
BogoMIPS:                        4608.00
Flags:                           fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc art arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc cpuid aperfmperf tsc_known_freq pni pclmulqdq dtes64 monitor ds_cpl vmx est tm2 ssse3 sdbg fma cx16 xtpr pdcm pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch cpuid_fault epb cat_l2 invpcid_single cdp_l2 ssbd ibrs ibpb stibp ibrs_enhanced tpr_shadow vnmi flexpriority ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid rdt_a avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb intel_pt avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves split_lock_detect dtherm ida arat pln pts hwp hwp_notify hwp_act_window hwp_epp hwp_pkg_req avx512vbmi umip pku ospke avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg avx512_vpopcntdq rdpid movdiri movdir64b fsrm avx512_vp2intersect md_clear ibt flush_l1d arch_capabilities
Virtualization:                  VT-x
L1d cache:                       384 KiB (8 instances)
L1i cache:                       256 KiB (8 instances)
L2 cache:                        10 MiB (8 instances)
L3 cache:                        24 MiB (1 instance)
NUMA node(s):                    1
NUMA node0 CPU(s):               0-15
Vulnerability Itlb multihit:     Not affected
Vulnerability L1tf:              Not affected
Vulnerability Mds:               Not affected
Vulnerability Meltdown:          Not affected
Vulnerability Mmio stale data:   Not affected
Vulnerability Retbleed:          Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:        Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:        Mitigation; Enhanced IBRS, IBPB conditional, RSB filling, PBRSB-eIBRS SW sequence
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Not affected

Versions of relevant libraries:
[pip3] numpy==1.23.0
[pip3] torch==2.0.1
[pip3] torchmetrics==0.8.0
[pip3] torchvision==0.15.2
[pip3] triton==2.0.0
[conda] Could not collect

[BloodAxe]

Hi there.

Regarding YoloNAS + segmentation I don't have an answer for you at the moment whether will or will not it will be added. In theory one can build such model relatively easy, however I doubt it would be that accurate/efficient compared to models purposely built for low-latency image segmentation.

On the issues you are having - the choice of BinaryIOU / BinarySegmentationVisualizationCallback metric / visualization callback looks wrong to me. As name suggest these classes defined for binary segmentation task (single class). In case of COCO you are dealing with multiclass segmentation and therefore should be using metric class designed for that task: IoU

[23pointsNorth]

Previous wrong assumption

You are right, didn't catch the `binaryIOU` part. Also, given that `BinarySegmentationVisualizationCallback` exists, why does `SegmentationVisualizationCallback` doesn't exist? In essence looking around the recepies, I am having a hard time converting the default [COCO one](https://github.com/Deci-AI/super-gradients/blob/d822731c8d0a06a649265c9796063e1fb1e8da26/src/super_gradients/recipes/training_hyperparams/coco2017_yolo_nas_train_params.yaml#L21) to a training params dict. How would that look like? Minimizing the amount of additional params to e.g. ``` train_params = {"max_epochs": 30, "lr_mode": "cosine", "initial_lr": 0.01, "optimizer": "AdamW", "loss": PPYoloELoss(num_classes=len(train_loader.dataset.dataset_classes_inclusion_tuples_list)), "ema": True, "zero_weight_decay_on_bias_and_bn": True, "loss_logging_items_names": ["loss"], } ``` Still results in ``` File "/home/bot/.local/lib/python3.10/site-packages/super_gradients/training/sg_trainer/sg_trainer.py", line 475, in _get_losses loss = self.criterion(outputs, targets) File "/home/bot/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl return forward_call(*args, **kwargs) File "/home/bot/.local/lib/python3.10/site-packages/super_gradients/training/losses/ppyolo_loss.py", line 744, in forward ( ValueError: too many values to unpack (expected 6) ```

For future reference for multiclass segmentation with COCO dataset this may be a good starting point.

import os
from super_gradients import Trainer, setup_device
from super_gradients.training import MultiGPUMode, dataloaders, models
from super_gradients.common.object_names import (
    Losses,
    LRWarmups,
    LRSchedulers,
    Optimizers,
    Metrics,
    Models,
)

from super_gradients.training.metrics.segmentation_metrics import IoU
from super_gradients.training.utils.callbacks import (
    BinarySegmentationVisualizationCallback,
    Phase,
)

BATCH_SIZE = 8
EPOCHS = 5
CHECKPOINT_DIR = "ckpts/"

os.makedirs(CHECKPOINT_DIR, exist_ok=True)
setup_device(device="cuda", multi_gpu=MultiGPUMode.AUTO)
trainer = Trainer(
    experiment_name="segmentation_quick_start", ckpt_root_dir=CHECKPOINT_DIR
)

train_loader = dataloaders.coco_segmentation_train(
    dataset_params={"root_dir": "data/coco"},
    dataloader_params={"batch_size": BATCH_SIZE},
)
valid_loader = dataloaders.coco_segmentation_val(
    dataset_params={"root_dir": "data/coco"},
    dataloader_params={"batch_size": BATCH_SIZE},
)

print("Dataloader parameters:", train_loader.dataloader_params)
print("Dataset parameters", train_loader.dataset.dataset_params)

model = models.get(
    model_name=Models.PP_LITE_T_SEG,
    num_classes=len(train_loader.dataset.classes),
)

train_params = {
    # ENABLING SILENT MODE
    # "silent_mode": True,
    "average_best_models": True,
    "warmup_mode": LRWarmups.LINEAR_EPOCH_STEP,
    "warmup_initial_lr": 1e-6,
    "lr_warmup_epochs": 3,
    "initial_lr": 5e-4,
    "lr_mode": LRSchedulers.COSINE,
    "cosine_final_lr_ratio": 0.1,
    "optimizer": Optimizers.ADAMW,
    "optimizer_params": {"weight_decay": 0.0001},
    "zero_weight_decay_on_bias_and_bn": True,
    "ema": True,
    "ema_params": {"decay": 0.9, "decay_type": "threshold"},
    "max_epochs": EPOCHS,
    # "mixed_precision": True,  # Allow for CPU/GPU testing
    "loss": Losses.CROSS_ENTROPY,
    "train_metrics_list": [IoU(num_classes=len(train_loader.dataset.classes))],
    "valid_metrics_list": [IoU(num_classes=len(train_loader.dataset.classes))],
    "metric_to_watch": Metrics.IOU,
    "phase_callbacks": [
        BinarySegmentationVisualizationCallback(
            phase=Phase.VALIDATION_BATCH_END, freq=1, last_img_idx_in_batch=4
        )
    ],
    "loss_logging_items_names": ["loss"],
}

trainer.train(
    model=model,
    training_params=train_params,
    train_loader=train_loader,
    valid_loader=valid_loader,
)

[vivekbharadhwajsa]

The COCO Segmentation dataloader always redirects to CoCoSegmentationDataSet and results in the error : 'CoCoSegmentationDataSet' object has no attribute 'samples_dir_suffix' .

Using the CoCoSegmentationDataSet to create a dataset will also always result in the same error.

The documentation in the code for CoCoSegmentationDataSet is also wrong as it says the usage as : but the init for CoCoSegmentationDataSet requires root_dir.

Could you please suggest a working example for COCO Segmentation dataset?