This commit is contained in:
Glenn Jocher 2019-08-21 01:35:29 +02:00
parent 9c6b6968ba
commit 1ec9527f24
1 changed files with 7 additions and 3 deletions

View File

@ -4,6 +4,7 @@ from utils.parse_config import *
from utils.utils import * from utils.utils import *
ONNX_EXPORT = False ONNX_EXPORT = False
arc = 'normal' # (normal, uCE, uBCE, uBCEs) detection architectures
def create_modules(module_defs, img_size): def create_modules(module_defs, img_size):
@ -77,8 +78,12 @@ def create_modules(module_defs, img_size):
# Initialize preceding Conv2d() bias (https://arxiv.org/pdf/1708.02002.pdf section 3.3) # Initialize preceding Conv2d() bias (https://arxiv.org/pdf/1708.02002.pdf section 3.3)
bias = module_list[-1][0].bias.view(len(mask), -1) # 255 to 3x85 bias = module_list[-1][0].bias.view(len(mask), -1) # 255 to 3x85
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':
bias[:, 4] += 3.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
@ -168,7 +173,6 @@ class YOLOLayer(nn.Module):
# io[..., 2:4] = ((torch.sigmoid(io[..., 2:4]) * 2) ** 3) * self.anchor_wh # wh power method # io[..., 2:4] = ((torch.sigmoid(io[..., 2:4]) * 2) ** 3) * self.anchor_wh # wh power method
io[..., :4] *= self.stride io[..., :4] *= self.stride
arc = 'normal' # (normal, uCE, uBCE, uBCEs) detection architectures
if arc == 'normal': if arc == 'normal':
torch.sigmoid_(io[..., 4:]) torch.sigmoid_(io[..., 4:])
elif arc == 'uCE': # unified CE (1 background + 80 classes) elif arc == 'uCE': # unified CE (1 background + 80 classes)