[ConvNets/PyT] TorchHub improvements

Author: asulecki

PyTorch/Classification/ConvNets/image_classification/models/__init__.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .entrypoints import nvidia_efficientnet, nvidia_resneXt, nvidia_resnet50, nvidia_convnets_processing_utils
+from .entrypoints import nvidia_convnets_processing_utils, nvidia_efficientnet
 from .resnet import resnet50, resnext101_32x4d, se_resnext101_32x4d
 from .efficientnet import (
     efficientnet_b0,

PyTorch/Classification/ConvNets/image_classification/models/efficientnet.py

@@ -573,7 +573,7 @@ def _m(*args, **kwargs):
 
 # }}}
 
-_ce = lambda n: EntryPoint(n, architectures[n])
+_ce = lambda n: EntryPoint.create(n, architectures[n])
 efficientnet_b0 = _ce("efficientnet-b0")
 efficientnet_b4 = _ce("efficientnet-b4")
 

PyTorch/Classification/ConvNets/image_classification/models/entrypoints.py

@@ -27,18 +27,6 @@
 # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
-def nvidia_resnet50(pretrained=True, **kwargs):
-    """Constructs a ResNet50 model.
-    For detailed information on model input and output, training recipies, inference and performance
-    visit: github.com/NVIDIA/DeepLearningExamples and/or ngc.nvidia.com
-    Args:
-        pretrained (bool, True): If True, returns a model pretrained on IMAGENET dataset.
-    """
-
-    from . import resnet50
-
-    return resnet50(pretrained=pretrained)
-
 def nvidia_efficientnet(type='efficient-b0', pretrained=True, **kwargs):
     """Constructs a EfficientNet model.
     For detailed information on model input and output, training recipies, inference and performance
@@ -51,17 +39,6 @@ def nvidia_efficientnet(type='efficient-b0', pretrained=True, **kwargs):
 
     return _ce(type)(pretrained=pretrained, **kwargs)
 
-def nvidia_resneXt(pretrained=True, **kwargs):
-    """Constructs a ResNeXt model.
-    For detailed information on model input and output, training recipies, inference and performance
-    visit: github.com/NVIDIA/DeepLearningExamples and/or ngc.nvidia.com
-    Args:
-        pretrained (bool, True): If True, returns a model pretrained on IMAGENET dataset.
-    """
-
-    from . import resnext101_32x4d
-
-    return resnext101_32x4d(pretrained=pretrained)
 
 def nvidia_convnets_processing_utils():
     import numpy as np
@@ -135,4 +112,4 @@ def get_imgnet_classes():
 
             return imgnet_classes
 
-    return Processing()
+    return Processing()

PyTorch/Classification/ConvNets/image_classification/models/model.py

@@ -39,14 +39,28 @@ class Model:
     checkpoint_url: Optional[str] = None
 
 
+def torchhub_docstring(name: str):
+    return f"""Constructs a {name} model.
+    For detailed information on model input and output, training recipies, inference and performance
+    visit: github.com/NVIDIA/DeepLearningExamples and/or ngc.nvidia.com
+    Args:
+        pretrained (bool, True): If True, returns a model pretrained on IMAGENET dataset.
+    """
+
 class EntryPoint:
+    @staticmethod
+    def create(name: str, model: Model):
+        ep = EntryPoint(name, model)
+        ep.__doc__ = torchhub_docstring(name)
+        return ep
+
     def __init__(self, name: str, model: Model):
         self.name = name
         self.model = model
 
     def __call__(
         self,
-        pretrained=False,
+        pretrained=True,
         pretrained_from_file=None,
         state_dict_key_map_fn=None,
         **kwargs,
@@ -60,7 +74,9 @@ def __call__(
         if pretrained:
             assert self.model.checkpoint_url is not None
             state_dict = torch.hub.load_state_dict_from_url(
-                self.model.checkpoint_url, map_location=torch.device("cpu"), progress=True
+                self.model.checkpoint_url,
+                map_location=torch.device("cpu"),
+                progress=True,
             )
 
         if pretrained_from_file is not None:

PyTorch/Classification/ConvNets/image_classification/models/resnet.py

@@ -450,7 +450,7 @@ def no_remap(s):
     ),
 }
 
-_ce = lambda n: EntryPoint(n, __models[n])
+_ce = lambda n: EntryPoint.create(n, __models[n])
 resnet50 = _ce("resnet50")
 resnext101_32x4d = _ce("resnext101-32x4d")
 se_resnext101_32x4d = _ce("se-resnext101-32x4d")

hubconf.py

@@ -4,7 +4,17 @@
 from PyTorch.Detection.SSD.ssd import nvidia_ssd, nvidia_ssd_processing_utils
 sys.path.append(os.path.join(sys.path[0], 'PyTorch/Detection/SSD'))
 
-from PyTorch.Classification.ConvNets.image_classification.models import nvidia_efficientnet, nvidia_resneXt, nvidia_resnet50, nvidia_convnets_processing_utils
+from PyTorch.Classification.ConvNets.image_classification.models import resnet50 as nvidia_resnet50
+from PyTorch.Classification.ConvNets.image_classification.models import resnext101_32x4d as nvidia_resnext101_32x4d
+from PyTorch.Classification.ConvNets.image_classification.models import se_resnext101_32x4d as nvidia_se_resnext101_32x4d
+from PyTorch.Classification.ConvNets.image_classification.models import efficientnet_b0 as nvidia_efficientnet_b0
+from PyTorch.Classification.ConvNets.image_classification.models import efficientnet_b4 as nvidia_efficientnet_b4
+from PyTorch.Classification.ConvNets.image_classification.models import efficientnet_widese_b0 as nvidia_efficientnet_widese_b0
+from PyTorch.Classification.ConvNets.image_classification.models import efficientnet_widese_b4 as nvidia_efficientnet_widese_b4
+from PyTorch.Classification.ConvNets.image_classification.models import nvidia_convnets_processing_utils
+
+from PyTorch.Classification.ConvNets.image_classification.models import resnext101_32x4d as nvidia_resneXt
+from PyTorch.Classification.ConvNets.image_classification.models import nvidia_efficientnet
 sys.path.append(os.path.join(sys.path[0], 'PyTorch/Classification/ConvNets/image_classification'))
 
 from PyTorch.SpeechSynthesis.Tacotron2.tacotron2 import nvidia_tacotron2