updates
This commit is contained in:
parent
0a04eb9ff1
commit
e27b124828
|
@ -194,7 +194,7 @@ class YOLOLayer(nn.Module):
|
|||
xy = torch.sigmoid(p[..., 0:2]) + grid_xy[0] # x, y
|
||||
wh = torch.exp(p[..., 2:4]) * anchor_wh[0] # width, height
|
||||
p_conf = torch.sigmoid(p[:, 4:5]) # Conf
|
||||
p_cls = F.softmax(p[:, 5:85], 1) * p_conf # SSD-like conf
|
||||
p_cls = F.softmax(p[:, 5:5 + self.nc], 1) * p_conf # SSD-like conf
|
||||
return torch.cat((xy / ngu[0], wh, p_conf, p_cls), 1).t()
|
||||
|
||||
# p = p.view(1, m, 5 + self.nc)
|
||||
|
|
Loading…
Reference in New Issue