Spedon
commited on
Commit
·
26cc346
1
Parent(s):
0c01441
Optimize decode_outputs for OpenVINO (#1535)
Browse filesAvoid ScatterND ops that will cause error in openvino model optimizer.
yolox/models/yolo_head.py
CHANGED
@@ -246,8 +246,7 @@ class YOLOXHead(nn.Module):
|
|
246 |
grids = torch.cat(grids, dim=1).type(dtype)
|
247 |
strides = torch.cat(strides, dim=1).type(dtype)
|
248 |
|
249 |
-
outputs
|
250 |
-
outputs[..., 2:4] = torch.exp(outputs[..., 2:4]) * strides
|
251 |
return outputs
|
252 |
|
253 |
def get_losses(
|
|
|
246 |
grids = torch.cat(grids, dim=1).type(dtype)
|
247 |
strides = torch.cat(strides, dim=1).type(dtype)
|
248 |
|
249 |
+
outputs = torch.cat([(outputs[..., 0:2] + grids) * strides, torch.exp(outputs[..., 2:4]) * strides, outputs[..., 4:]], dim=-1)
|
|
|
250 |
return outputs
|
251 |
|
252 |
def get_losses(
|