Merge remote-tracking branch 'origin/master'
This commit is contained in:
commit
267b4301e2
|
@ -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
|
||||
|
||||
|
|
Loading…
Reference in New Issue