zhang chen
3 years ago
1 changed files with
1 additions and
0 deletions
-
pytorch/model.py
|
@ -41,6 +41,7 @@ class Model(): |
|
|
self._model.eval() |
|
|
self._model.eval() |
|
|
|
|
|
|
|
|
def __call__(self, img_tensor: torch.Tensor): |
|
|
def __call__(self, img_tensor: torch.Tensor): |
|
|
|
|
|
self._model.eval() |
|
|
features = self._model.forward_features(img_tensor) |
|
|
features = self._model.forward_features(img_tensor) |
|
|
if features.dim() == 4: # if the shape of feature map is [N, C, H, W], where H > 1 and W > 1 |
|
|
if features.dim() == 4: # if the shape of feature map is [N, C, H, W], where H > 1 and W > 1 |
|
|
global_pool = nn.AdaptiveAvgPool2d(1) |
|
|
global_pool = nn.AdaptiveAvgPool2d(1) |
|
|