Add inline documentation for models

torchvision/models/__init__.py

@@ -1,3 +1,33 @@
+"""The models subpackage contains definitions for the following model
+architectures:
+
+-  `AlexNet`_
+-  `VGG`_
+-  `ResNet`_
+
+You can construct a model with random weights by calling its constructor:
+
+.. code:: python
+
+    import torchvision.models as models
+    resnet18 = models.resnet18()
+    alexnet = models.alexnet()
+
+We provide pre-trained models for the ResNet variants and AlexNet, using the
+PyTorch :mod:`torch.utils.model_zoo`. These can  constructed by passing
+``pretrained=True``:
+
+.. code:: python
+
+    import torchvision.models as models
+    resnet18 = models.resnet18(pretrained=True)
+    alexnet = models.alexnet(pretrained=True)
+
+.. _AlexNet: https://arxiv.org/abs/1404.5997
+.. _VGG: https://arxiv.org/abs/1409.1556
+.. _ResNet: https://arxiv.org/abs/1512.03385
+"""
+
 from .alexnet import *
 from .resnet import *
 from .vgg import *

torchvision/models/alexnet.py

@@ -46,8 +46,11 @@ def forward(self, x):
 
 
 def alexnet(pretrained=False):
-    r"""AlexNet model architecture from the "One weird trick" paper.
-    https://arxiv.org/abs/1404.5997
+    r"""AlexNet model architecture from the
+    `"One weird trick..." <https://arxiv.org/abs/1404.5997>`_ paper.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
     """
     model = AlexNet()
     if pretrained:

torchvision/models/resnet.py

@@ -12,6 +12,7 @@
     'resnet34': 'https://s3.amazonaws.com/pytorch/models/resnet34-333f7ec4.pth',
     'resnet50': 'https://s3.amazonaws.com/pytorch/models/resnet50-19c8e357.pth',
     'resnet101': 'https://s3.amazonaws.com/pytorch/models/resnet101-5d3b4d8f.pth',
+    'resnet152': 'https://s3.amazonaws.com/pytorch/models/resnet152-b121ed2d.pth',
 }
 
 
@@ -152,32 +153,60 @@ def forward(self, x):
 
 
 def resnet18(pretrained=False):
+    """Constructs a ResNet-18 model.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
     model = ResNet(BasicBlock, [2, 2, 2, 2])
     if pretrained:
         model.load_state_dict(model_zoo.load_url(model_urls['resnet18']))
     return model
 
 
 def resnet34(pretrained=False):
+    """Constructs a ResNet-34 model.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
     model = ResNet(BasicBlock, [3, 4, 6, 3])
     if pretrained:
         model.load_state_dict(model_zoo.load_url(model_urls['resnet34']))
     return model
 
 
 def resnet50(pretrained=False):
+    """Constructs a ResNet-50 model.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
     model = ResNet(Bottleneck, [3, 4, 6, 3])
     if pretrained:
         model.load_state_dict(model_zoo.load_url(model_urls['resnet50']))
     return model
 
 
 def resnet101(pretrained=False):
+    """Constructs a ResNet-101 model.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
     model = ResNet(Bottleneck, [3, 4, 23, 3])
     if pretrained:
         model.load_state_dict(model_zoo.load_url(model_urls['resnet101']))
     return model
 
 
-def resnet152():
-    return ResNet(Bottleneck, [3, 8, 36, 3])
+def resnet152(pretrained=False):
+    """Constructs a ResNet-152 model.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = ResNet(Bottleneck, [3, 8, 36, 3])
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(model_urls['resnet152']))
+    return model

torchvision/models/vgg.py

@@ -53,32 +53,40 @@ def make_layers(cfg, batch_norm=False):
 
 
 def vgg11():
+    """VGG 11-layer model (configuration "A")"""
     return VGG(make_layers(cfg['A']))
 
 
 def vgg11_bn():
+    """VGG 11-layer model (configuration "A") with batch normalization"""
     return VGG(make_layers(cfg['A'], batch_norm=True))
 
 
 def vgg13():
+    """VGG 13-layer model (configuration "B")"""
     return VGG(make_layers(cfg['B']))
 
 
 def vgg13_bn():
+    """VGG 13-layer model (configuration "B") with batch normalization"""
     return VGG(make_layers(cfg['B'], batch_norm=True))
 
 
 def vgg16():
+    """VGG 11-layer model (configuration "B")"""
     return VGG(make_layers(cfg['D']))
 
 
 def vgg16_bn():
+    """VGG 16-layer model (configuration "D") with batch normalization"""
     return VGG(make_layers(cfg['D'], batch_norm=True))
 
 
 def vgg19():
+    """VGG 19-layer model (configuration "D")"""
     return VGG(make_layers(cfg['E']))
 
 
 def vgg19_bn():
+    """VGG 19-layer model (configuration 'E') with batch normalization"""
     return VGG(make_layers(cfg['E'], batch_norm=True))