Merge pull request #250 from rvandeghen/patch-1

lukemelas · web-flow · commit 7e8b0d312162 · 2021-04-15T11:16:36.000-04:00
Add new checkpoint
diff --git a/efficientnet_pytorch/model.py b/efficientnet_pytorch/model.py
@@ -238,18 +238,18 @@ def extract_endpoints(self, inputs):
         Returns:
             Dictionary of last intermediate features
             with reduction levels i in [1, 2, 3, 4, 5].
-
-        Example:
-            >>> import torch
-            >>> from efficientnet.model import EfficientNet
-            >>> inputs = torch.rand(1, 3, 224, 224)
-            >>> model = EfficientNet.from_pretrained('efficientnet-b0')
-            >>> endpoints = model.extract_endpoints(inputs)
-            >>> print(endpoints['reduction_1'].shape)  # torch.Size([1, 16, 112, 112])
-            >>> print(endpoints['reduction_2'].shape)  # torch.Size([1, 24, 56, 56])
-            >>> print(endpoints['reduction_3'].shape)  # torch.Size([1, 40, 28, 28])
-            >>> print(endpoints['reduction_4'].shape)  # torch.Size([1, 112, 14, 14])
-            >>> print(endpoints['reduction_5'].shape)  # torch.Size([1, 1280, 7, 7])
+            Example:
+                >>> import torch
+                >>> from efficientnet.model import EfficientNet
+                >>> inputs = torch.rand(1, 3, 224, 224)
+                >>> model = EfficientNet.from_pretrained('efficientnet-b0')
+                >>> endpoints = model.extract_endpoints(inputs)
+                >>> print(endpoints['reduction_1'].shape)  # torch.Size([1, 16, 112, 112])
+                >>> print(endpoints['reduction_2'].shape)  # torch.Size([1, 24, 56, 56])
+                >>> print(endpoints['reduction_3'].shape)  # torch.Size([1, 40, 28, 28])
+                >>> print(endpoints['reduction_4'].shape)  # torch.Size([1, 112, 14, 14])
+                >>> print(endpoints['reduction_5'].shape)  # torch.Size([1, 320, 7, 7])
+                >>> print(endpoints['reduction_6'].shape)  # torch.Size([1, 1280, 7, 7])
         """
         endpoints = dict()
 
@@ -265,6 +265,8 @@ def extract_endpoints(self, inputs):
             x = block(x, drop_connect_rate=drop_connect_rate)
             if prev_x.size(2) > x.size(2):
                 endpoints['reduction_{}'.format(len(endpoints) + 1)] = prev_x
+            elif idx == len(self._blocks) - 1:
+                endpoints['reduction_{}'.format(len(endpoints) + 1)] = x
             prev_x = x
 
         # Head
