updates

2019-08-22 15:48:06 +02:00 · 2019-08-22 15:48:06 +02:00 · 11aac8930b
parent baf6188df5
commit 11aac8930b
1 changed files with 4 additions and 1 deletions
--- a/models.py
+++ b/models.py
@ -81,9 +81,12 @@ def create_modules(module_defs, img_size):
            if arc == 'normal':
                bias[:, 4] -= 5.0  # obj
                bias[:, 5:] -= 4.0  # cls
-            elif arc == 'uCE':
+            elif arc == 'uCE':  # unified CE (1 background + 80 classes)
                bias[:, 4] += 3.0  # obj
                bias[:, 5:] -= 4.0  # cls
+            elif arc == 'uBCE':  # unified BCE (80 classes)
+                bias[:, 4] -= 5.0  # obj
+                bias[:, 5:] -= 4.0  # cls
            module_list[-1][0].bias = torch.nn.Parameter(bias.view(-1))

            # for l in model.yolo_layers:  # print pretrained biases