Cleanup- edit for Readability (#417)

Cleanup- edit for Readability
This commit is contained in:
Yonghye Kwon 2019-08-02 21:49:08 +09:00 committed by Glenn Jocher
parent 8501981f09
commit 2a284aacd0
1 changed files with 3 additions and 1 deletions

View File

@ -156,10 +156,12 @@ 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[..., :4] *= self.stride
# 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
if self.nc == 1:
io[..., 5] = 1 # single-class model https://github.com/ultralytics/yolov3/issues/235