From d4c8c4fa4b6bef8d2a837dcd5ea938e43b73e8cf Mon Sep 17 00:00:00 2001
From: Phat Tran <36766404+ptran1203@users.noreply.github.com>
Date: Mon, 29 Mar 2021 23:45:46 +0700
Subject: [PATCH] Add --label-smoothing eps argument to train.py (default 0.0)
 (#2344)

* Add label smoothing option

* Correct data type

* add_log

* Remove log

* Add log

* Update loss.py

remove comment (too versbose)

Co-authored-by: phattran <phat.tranhoang@cyberlogitec.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
---
 train.py      | 2 ++
 utils/loss.py | 2 +-
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/train.py b/train.py
index d5b2d1b75c52..d55833bf45a3 100644
--- a/train.py
+++ b/train.py
@@ -224,6 +224,7 @@ def train(hyp, opt, device, tb_writer=None):
     hyp['box'] *= 3. / nl  # scale to layers
     hyp['cls'] *= nc / 80. * 3. / nl  # scale to classes and layers
     hyp['obj'] *= (imgsz / 640) ** 2 * 3. / nl  # scale to image size and layers
+    hyp['label_smoothing'] = opt.label_smoothing
     model.nc = nc  # attach number of classes to model
     model.hyp = hyp  # attach hyperparameters to model
     model.gr = 1.0  # iou loss ratio (obj_loss = 1.0 or iou)
@@ -481,6 +482,7 @@ def train(hyp, opt, device, tb_writer=None):
     parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
     parser.add_argument('--quad', action='store_true', help='quad dataloader')
     parser.add_argument('--linear-lr', action='store_true', help='linear LR')
+    parser.add_argument('--label-smoothing', type=float, default=0.0, help='Label smoothing epsilon')
     parser.add_argument('--upload_dataset', action='store_true', help='Upload dataset as W&B artifact table')
     parser.add_argument('--bbox_interval', type=int, default=-1, help='Set bounding-box image logging interval for W&B')
     parser.add_argument('--save_period', type=int, default=-1, help='Log model after every "save_period" epoch')
diff --git a/utils/loss.py b/utils/loss.py
index 2302d18de87d..9e78df17fdf3 100644
--- a/utils/loss.py
+++ b/utils/loss.py
@@ -97,7 +97,7 @@ def __init__(self, model, autobalance=False):
         BCEobj = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['obj_pw']], device=device))
 
         # Class label smoothing https://arxiv.org/pdf/1902.04103.pdf eqn 3
-        self.cp, self.cn = smooth_BCE(eps=0.0)
+        self.cp, self.cn = smooth_BCE(eps=h.get('label_smoothing', 0.0))  # positive, negative BCE targets
 
         # Focal loss
         g = h['fl_gamma']  # focal loss gamma