Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add support for torchvision's mobilenet_v2 #2456

Merged
merged 1 commit into from Jan 15, 2020
Merged

Add support for torchvision's mobilenet_v2 #2456

merged 1 commit into from Jan 15, 2020

Conversation

rsomani95
Copy link
Contributor

This PR enables you to use the following code:

from fastai.vision import *

data  = #...
arch  = models.mobilenet_v2
learn = cnn_learner(data, arch, pretrained=True)

Earlier, you would have to do something this instead:

import torchvision

mobilenet_split = lambda m: (m[0][0][10], m[1])
arch  = torchvision.models.mobilenet_v2
learn = cnn_learner(data, arch, cut=-1, split_on=mobilenet_split)

There is a brief discussion here, but it isn't a dedicated thread.
I took the liberty of making a PR anyways, hope that's alright.

rsomani95
rsomani95 commented Jan 15, 2020
View reviewed changes
fastai/vision/learner.py
@@ -19,13 +19,15 @@ def _squeezenet_split(m:nn.Module): return (m[0][0][5], m[0][0][8], m[1])
def _densenet_split(m:nn.Module): return (m[0][0][7],m[1])
def _vgg_split(m:nn.Module): return (m[0][0][22],m[1])
def _alexnet_split(m:nn.Module): return (m[0][0][6],m[1])
def _mobilenetv2_split(m:nn.Module): return (m[0][0][10],m[1])
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This was a somewhat arbitrary choice, but it splits the network into roughly equal parts

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sounds fine to me!

@sgugger sgugger merged commit 67d951c into fastai:master Jan 15, 2020
@sgugger
Copy link
Contributor

sgugger commented Jan 15, 2020

Thanks a lot!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@sgugger sgugger sgugger left review comments

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.

None yet
2 participants
@rsomani95 @sgugger