vision
vision copied to clipboard
pytorch
add lovasz hinge loss optional to segmentation loss
🚀 The feature
Hello! In references/segmentation/train.py
def criterion(inputs, target): <- 56th line
losses = {}
for name, x in inputs.items():
losses[name] = nn.functional.cross_entropy(x, target, ignore_index=255)
if len(losses) == 1:
return losses["out"]
return losses["out"] + 0.5 * losses["aux"]
It is only binary cross entropy. But as far as I Know many papers use lovasz hinge loss to enhance model's ability to segment better in edge, boundary which is uncertain.
What about adding this loss to model optional.
environment :
python : 3.10
NVIDIA GeForce GTX 1650 GPU
experiment script
from ultralytics import YOLO
import matplotlib.pyplot as plt
import cv2
import torch
import os
def main():
model = YOLO("yolov8s-seg.pt")
results = model.train(
data="coco128-seg.yaml",
epochs=5,
imgsz=640,
batch=8,
workers=2,
name="yolov8-segmentation-exp",
device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
)
results_dir = "runs/segment/yolov8-segmentation-exp/"
results_img_path = os.path.join(results_dir, "results.png")
if os.path.exists(results_img_path):
img = cv2.imread(results_img_path)
img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
plt.figure(figsize=(10, 6))
plt.imshow(img_rgb)
plt.axis('off')
plt.title("YOLOv8 Training Metrics")
plt.show()
if __name__ == '__main__':
from multiprocessing import freeze_support
freeze_support()
main()
changed version(lovasz_weight : 0.008)
Starting training for 5 epochs...
Epoch GPU_mem box_loss seg_loss cls_loss dfl_loss Instances Size
1/5 5.23G 0.9426 2.309 1.233 1.093 82 640: 100%|██████████| 16/16 [02:10<00:00, 8.15s/it]
Class Images Instances Box(P R mAP50 mAP50-95) Mask(P R mAP50 mAP50-95): 100%|██████████| 8/8 [00:47<00:00, 5.97s/it]
all 128 929 0.755 0.679 0.761 0.594 0.742 0.652 0.725 0.489
0%| | 0/16 [00:00<?, ?it/s]
Epoch GPU_mem box_loss seg_loss cls_loss dfl_loss Instances Size
2/5 5.42G 1.028 2.44 1.172 1.111 89 640: 100%|██████████| 16/16 [01:52<00:00, 7.02s/it]
Class Images Instances Box(P R mAP50 mAP50-95) Mask(P R mAP50 mAP50-95): 100%|██████████| 8/8 [00:29<00:00, 3.74s/it]
all 128 929 0.77 0.698 0.776 0.605 0.743 0.674 0.729 0.494
Epoch GPU_mem box_loss seg_loss cls_loss dfl_loss Instances Size
3/5 5.34G 0.9623 2.313 1.044 1.068 128 640: 100%|██████████| 16/16 [01:57<00:00, 7.35s/it]
Class Images Instances Box(P R mAP50 mAP50-95) Mask(P R mAP50 mAP50-95): 100%|██████████| 8/8 [00:52<00:00, 6.62s/it]
all 128 929 0.809 0.7 0.787 0.612 0.784 0.677 0.742 0.5
Epoch GPU_mem box_loss seg_loss cls_loss dfl_loss Instances Size
4/5 5.78G 0.9621 2.282 1.112 1.106 39 640: 100%|██████████| 16/16 [02:14<00:00, 8.38s/it]
Class Images Instances Box(P R mAP50 mAP50-95) Mask(P R mAP50 mAP50-95): 100%|██████████| 8/8 [00:27<00:00, 3.41s/it]
all 128 929 0.835 0.696 0.79 0.617 0.809 0.668 0.745 0.505
0%| | 0/16 [00:00<?, ?it/s]
Epoch GPU_mem box_loss seg_loss cls_loss dfl_loss Instances Size
5/5 5.43G 0.997 2.383 1.021 1.073 85 640: 100%|██████████| 16/16 [01:58<00:00, 7.39s/it]
Class Images Instances Box(P R mAP50 mAP50-95) Mask(P R mAP50 mAP50-95): 100%|██████████| 8/8 [00:27<00:00, 3.43s/it]
all 128 929 0.816 0.72 0.797 0.624 0.795 0.678 0.746 0.509
original version
original version
Starting training for 5 epochs...
Epoch GPU_mem box_loss seg_loss cls_loss dfl_loss Instances Size
1/5 4.58G 0.9426 2.157 1.233 1.093 82 640: 100%|██████████| 16/16 [01:16<00:00, 4.79s/it]
Class Images Instances Box(P R mAP50 mAP50-95) Mask(P R mAP50 mAP50-95): 100%|██████████| 8/8 [00:30<00:00, 3.86s/it]
all 128 929 0.755 0.679 0.761 0.594 0.743 0.653 0.725 0.489
Epoch GPU_mem box_loss seg_loss cls_loss dfl_loss Instances Size
2/5 4.74G 1.029 2.292 1.171 1.111 89 640: 100%|██████████| 16/16 [01:08<00:00, 4.29s/it]
Class Images Instances Box(P R mAP50 mAP50-95) Mask(P R mAP50 mAP50-95): 100%|██████████| 8/8 [00:32<00:00, 4.09s/it]
all 128 929 0.769 0.699 0.775 0.604 0.741 0.674 0.728 0.495
0%| | 0/16 [00:00<?, ?it/s]
Epoch GPU_mem box_loss seg_loss cls_loss dfl_loss Instances Size
3/5 4.67G 0.9636 2.168 1.043 1.068 128 640: 100%|██████████| 16/16 [01:04<00:00, 4.06s/it]
Class Images Instances Box(P R mAP50 mAP50-95) Mask(P R mAP50 mAP50-95): 100%|██████████| 8/8 [00:31<00:00, 3.89s/it]
all 128 929 0.808 0.699 0.788 0.614 0.781 0.671 0.741 0.5
Epoch GPU_mem box_loss seg_loss cls_loss dfl_loss Instances Size
4/5 4.87G 0.963 2.135 1.112 1.107 39 640: 100%|██████████| 16/16 [01:08<00:00, 4.26s/it]
Class Images Instances Box(P R mAP50 mAP50-95) Mask(P R mAP50 mAP50-95): 100%|██████████| 8/8 [00:36<00:00, 4.61s/it]
all 128 929 0.829 0.699 0.791 0.617 0.809 0.667 0.745 0.502
Epoch GPU_mem box_loss seg_loss cls_loss dfl_loss Instances Size
5/5 4.7G 0.9968 2.242 1.019 1.073 85 640: 100%|██████████| 16/16 [00:52<00:00, 3.30s/it]
Class Images Instances Box(P R mAP50 mAP50-95) Mask(P R mAP50 mAP50-95): 100%|██████████| 8/8 [00:34<00:00, 4.29s/it]
all 128 929 0.816 0.72 0.797 0.625 0.79 0.687 0.75 0.507
I didn't do such as lovasz_loss * min(epoch_num, 10)*0.1.
if lovasz weight is big, test result was not good.
Motivation, pitch
Bce loss is well generalized loss but adding lovasz when it is needed can enhance model's ability.
Alternatives
No response
Additional context
No response
