updates
This commit is contained in:
parent
a59caf053a
commit
efc662351b
12
models.py
12
models.py
|
@ -137,17 +137,17 @@ class YOLOLayer(nn.Module):
|
||||||
anchor_wh = self.anchor_wh.repeat((1, 1, nG, nG, 1)).view((1, -1, 2)) / nG
|
anchor_wh = self.anchor_wh.repeat((1, 1, nG, nG, 1)).view((1, -1, 2)) / nG
|
||||||
|
|
||||||
# p = p.view(-1, 5 + self.nC)
|
# p = p.view(-1, 5 + self.nC)
|
||||||
# xy = xy + self.grid_xy[0] # x, y
|
# xy = torch.sigmoid(p[..., 0:2]) + grid_xy[0] # x, y
|
||||||
# wh = torch.exp(wh) * self.anchor_wh[0] # width, height
|
# wh = torch.exp(p[..., 2:4]) * anchor_wh[0] # width, height
|
||||||
# p_conf = torch.sigmoid(p[:, 4:5]) # Conf
|
# p_conf = torch.sigmoid(p[:, 4:5]) # Conf
|
||||||
# p_cls = F.softmax(p[:, 5:], 1) * p_conf # SSD-like conf
|
# p_cls = F.softmax(p[:, 5:85], 1) * p_conf # SSD-like conf
|
||||||
# return torch.cat((xy / nG, wh, p_conf, p_cls), 1).t()
|
# return torch.cat((xy / nG, wh, p_conf, p_cls), 1).t()
|
||||||
|
|
||||||
p = p.view(1, -1, 5 + self.nC)
|
p = p.view(1, -1, 5 + self.nC)
|
||||||
xy = torch.sigmoid(p[..., 0:2]) + grid_xy # x, y
|
xy = torch.sigmoid(p[..., 0:2]) + grid_xy # x, y
|
||||||
wh = torch.exp(p[..., 2:4]) * anchor_wh # width, height
|
wh = torch.exp(p[..., 2:4]) * anchor_wh # width, height
|
||||||
p_conf = torch.sigmoid(p[..., 4:5]) # Conf
|
p_conf = torch.sigmoid(p[..., 4:5]) # Conf
|
||||||
p_cls = p[..., 5:]
|
p_cls = p[..., 5:85]
|
||||||
# Broadcasting only supported on first dimension in CoreML. See onnx-coreml/_operators.py
|
# Broadcasting only supported on first dimension in CoreML. See onnx-coreml/_operators.py
|
||||||
# p_cls = F.softmax(p_cls, 2) * p_conf # SSD-like conf
|
# p_cls = F.softmax(p_cls, 2) * p_conf # SSD-like conf
|
||||||
p_cls = torch.exp(p_cls).permute((2, 1, 0))
|
p_cls = torch.exp(p_cls).permute((2, 1, 0))
|
||||||
|
@ -203,8 +203,8 @@ class Darknet(nn.Module):
|
||||||
layer_outputs.append(x)
|
layer_outputs.append(x)
|
||||||
|
|
||||||
if ONNX_EXPORT:
|
if ONNX_EXPORT:
|
||||||
output = torch.cat(output, 1) # merge the 3 layers 85 x (507, 2028, 8112) to 85 x 10647
|
output = torch.cat(output, 1) # cat 3 layers 85 x (507, 2028, 8112) to 85 x 10647
|
||||||
return output[5:].t(), output[:4].t() # ONNX scores, boxes
|
return output[5:85].t(), output[:4].t() # ONNX scores, boxes
|
||||||
else:
|
else:
|
||||||
return output if self.training else torch.cat(output, 1)
|
return output if self.training else torch.cat(output, 1)
|
||||||
|
|
||||||
|
|
Loading…
Reference in New Issue