Browse Source
add omnivore
Signed-off-by: gexy5 <xinyu.ge@zilliz.com>
main
gexy5
2 years ago
1 changed files with
2 additions and
2 deletions
-
omnivore.py
|
|
@ -101,10 +101,10 @@ class Omnivore(NNOperator): |
|
|
) |
|
|
) |
|
|
inputs = data.to(self.device)[None, ...] |
|
|
inputs = data.to(self.device)[None, ...] |
|
|
|
|
|
|
|
|
feats = self.model.forward_features(inputs ,input_type = self.input_type) |
|
|
|
|
|
|
|
|
feats = self.model.forward_features(inputs, input_type = self.input_type) |
|
|
features = feats.to('cpu').squeeze(0).detach().numpy() |
|
|
features = feats.to('cpu').squeeze(0).detach().numpy() |
|
|
|
|
|
|
|
|
outs = self.model.head(feats) |
|
|
|
|
|
|
|
|
outs = self.model.head(feats, input_type = self.input_type) |
|
|
post_act = torch.nn.Softmax(dim=1) |
|
|
post_act = torch.nn.Softmax(dim=1) |
|
|
preds = post_act(outs) |
|
|
preds = post_act(outs) |
|
|
pred_scores, pred_classes = preds.topk(k=self.topk) |
|
|
pred_scores, pred_classes = preds.topk(k=self.topk) |
|
|
|
