This commit is contained in:
Glenn Jocher 2019-08-22 15:48:06 +02:00
parent baf6188df5
commit 11aac8930b
1 changed files with 4 additions and 1 deletions

View File

@ -81,9 +81,12 @@ def create_modules(module_defs, img_size):
if arc == 'normal': if arc == 'normal':
bias[:, 4] -= 5.0 # obj bias[:, 4] -= 5.0 # obj
bias[:, 5:] -= 4.0 # cls bias[:, 5:] -= 4.0 # cls
elif arc == 'uCE': elif arc == 'uCE': # unified CE (1 background + 80 classes)
bias[:, 4] += 3.0 # obj bias[:, 4] += 3.0 # obj
bias[:, 5:] -= 4.0 # cls 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)) module_list[-1][0].bias = torch.nn.Parameter(bias.view(-1))
# for l in model.yolo_layers: # print pretrained biases # for l in model.yolo_layers: # print pretrained biases