logo
Browse Source

add omnivore

Signed-off-by: gexy5 <xinyu.ge@zilliz.com>
main
gexy5 2 years ago
parent
commit
b31db9fd88
  1. 11
      omnivore.py

11
omnivore.py

@ -99,16 +99,7 @@ class Omnivore(NNOperator):
)
inputs = data.to(self.device)[None, ...]
feats = self.model.forward_features(inputs)
if self.model.reshape:
if self.model.is_shift and self.model.temporal_pool:
base_out = feats.view((-1, self.model.num_segments // 2) + feats.size()[1:])
else:
base_out = feats.view((-1, self.model.num_segments) + feats.size()[1:])
output = self.model.consensus(base_out)
features = output.to('cpu').squeeze(0).detach().numpy()
outs = self.model.head(feats)
outs = self.model.head(inputs)
post_act = torch.nn.Softmax(dim=1)
preds = post_act(outs)
pred_scores, pred_classes = preds.topk(k=self.topk)

Loading…
Cancel
Save