logo
Browse Source

add return shape introduction

main
zhang chen 4 years ago
parent
commit
27666263ee
  1. 1
      README.md
  2. 3
      pytorch/model.py

1
README.md

@ -41,6 +41,7 @@ The Operator returns a tuple `Tuple[('feature_vector', numpy.ndarray)]` containi
- feature_vector: - feature_vector:
- the embedding of the image - the embedding of the image
- data type: `numpy.ndarray` - data type: `numpy.ndarray`
- shape: (dim,)
## Requirements ## Requirements

3
pytorch/model.py

@ -32,8 +32,7 @@ class Model():
def __call__(self, img_tensor: torch.Tensor): def __call__(self, img_tensor: torch.Tensor):
features = self._model.extract_features(img_tensor) features = self._model.extract_features(img_tensor)
channels = features.shape[0]
return self._avg_pooling(features).view(channels, -1).detach().numpy()
return self._avg_pooling(features).flatten().detach().numpy()
def train(self): def train(self):
""" """

Loading…
Cancel
Save