Skip to content

Commit

Permalink
Add doc string for each available model in fastai library.
Browse files Browse the repository at this point in the history
The docs should help users to research details about the models.
  • Loading branch information
Piotr Czapla committed Mar 12, 2018
1 parent 4c2f4de commit 00a363b
Showing 1 changed file with 47 additions and 3 deletions.
50 changes: 47 additions & 3 deletions fastai/torch_imports.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,17 +31,61 @@ def load_pre(pre, f, fn):
if pre: load_model(m, f'{path}/weights/{fn}.pth')
return m

def inception_4(pre):
return children(load_pre(pre, InceptionV4, 'inceptionv4-97ef9c30'))[0]
def _fastai_model(name, paper_title, paper_href):
def add_docs_wrapper(f):
f.__doc__ = f"""{name} model from
`"{paper_title}" <{paper_href}>`_
Args:
pre (bool): If True, returns a model pre-trained on ImageNet
"""
return f
return add_docs_wrapper

@_fastai_model('Inception 4', 'Inception-v4, Inception-ResNet and the Impact of Residual Connections on Learning',
'https://arxiv.org/pdf/1602.07261.pdf')
def inception_4(pre): return children(load_pre(pre, InceptionV4, 'inceptionv4-97ef9c30'))[0]

@_fastai_model('Inception 4', 'Inception-v4, Inception-ResNet and the Impact of Residual Connections on Learning',
'https://arxiv.org/pdf/1602.07261.pdf')
def inceptionresnet_2(pre): return load_pre(pre, InceptionResnetV2, 'inceptionresnetv2-d579a627')

@_fastai_model('ResNeXt 50', 'Aggregated Residual Transformations for Deep Neural Networks',
'https://arxiv.org/abs/1611.05431')
def resnext50(pre): return load_pre(pre, resnext_50_32x4d, 'resnext_50_32x4d')

@_fastai_model('ResNeXt 101_32', 'Aggregated Residual Transformations for Deep Neural Networks',
'https://arxiv.org/abs/1611.05431')
def resnext101(pre): return load_pre(pre, resnext_101_32x4d, 'resnext_101_32x4d')

@_fastai_model('ResNeXt 101_64', 'Aggregated Residual Transformations for Deep Neural Networks',
'https://arxiv.org/abs/1611.05431')
def resnext101_64(pre): return load_pre(pre, resnext_101_64x4d, 'resnext_101_64x4d')

@_fastai_model('Inception 4', 'Wide Residual Networks',
'https://arxiv.org/pdf/1605.07146.pdf')
def wrn(pre): return load_pre(pre, wrn_50_2f, 'wrn_50_2f')

@_fastai_model('Densenet-121', 'Densely Connected Convolutional Networks',
'https://arxiv.org/pdf/1608.06993.pdf')
def dn121(pre): return children(densenet121(pre))[0]

@_fastai_model('Densenet-169', 'Densely Connected Convolutional Networks',
'https://arxiv.org/pdf/1608.06993.pdf')
def dn161(pre): return children(densenet161(pre))[0]

@_fastai_model('Densenet-161', 'Densely Connected Convolutional Networks',
'https://arxiv.org/pdf/1608.06993.pdf')
def dn169(pre): return children(densenet169(pre))[0]

@_fastai_model('Densenet-201', 'Densely Connected Convolutional Networks',
'https://arxiv.org/pdf/1608.06993.pdf')
def dn201(pre): return children(densenet201(pre))[0]

@_fastai_model('Vgg-16 with batch norm added', 'Very Deep Convolutional Networks for Large-Scale Image Recognition',
'https://arxiv.org/pdf/1409.1556.pdf')
def vgg16(pre): return children(vgg16_bn(pre))[0]
def vgg19(pre): return children(vgg19_bn(pre))[0]

@_fastai_model('Vgg-19 with batch norm added', 'Very Deep Convolutional Networks for Large-Scale Image Recognition',
'https://arxiv.org/pdf/1409.1556.pdf')
def vgg19(pre): return children(vgg19_bn(pre))[0]

0 comments on commit 00a363b

Please sign in to comment.