This commit is contained in:
Glenn Jocher 2019-12-04 17:50:52 -08:00
parent 0a04eb9ff1
commit e27b124828
1 changed files with 1 additions and 1 deletions

View File

@ -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)