This commit is contained in:
Glenn Jocher 2019-04-06 16:13:35 +02:00
parent 112f061f4e
commit 7ee48a43b6
1 changed files with 1 additions and 1 deletions

View File

@ -159,7 +159,7 @@ class YOLOLayer(nn.Module):
io = p.clone() # inference output
io[..., 0:2] = torch.sigmoid(io[..., 0:2]) + self.grid_xy # xy
io[..., 2:4] = torch.exp(io[..., 2:4]) * self.anchor_wh # wh yolo method
# io[..., 2:4] = ((torch.sigmoid(io[..., 2:4]) * 2) ** 2) * self.anchor_wh # wh power method
# io[..., 2:4] = ((torch.sigmoid(io[..., 2:4]) * 2) ** 3) * self.anchor_wh # wh power method
io[..., 4:] = torch.sigmoid(io[..., 4:]) # p_conf, p_cls
# io[..., 5:] = F.softmax(io[..., 5:], dim=4) # p_cls
io[..., :4] *= self.stride