Skip to content

make shufflenet and resnet scriptable #1270

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

Merged
merged 6 commits into from
Sep 2, 2019
Merged

make shufflenet and resnet scriptable #1270

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
bd3144c
make shufflenet scriptable
Aug 14, 2019
0f8a35d
make resnet18 scriptable
Aug 28, 2019
a298ed2
set downsample to identity instead of __constants__ api
Aug 28, 2019
a245b06
use __constants__ for downsample instead of identity
Aug 28, 2019
9ce4930
import tensor to fix flake
Aug 30, 2019
e14d72d
use torch.Tensor type annotation instead of import
Aug 30, 2019
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions test/test_models.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,9 +33,9 @@ def get_available_video_models():
"fcn_resnet101": False,
"googlenet": False,
"densenet121": False,
"resnet18": False,
"resnet18": True,
"alexnet": True,
"shufflenet_v2_x1_0": False,
"shufflenet_v2_x1_0": True,
"squeezenet1_0": True,
"vgg11": True,
"inception_v3": False,
Expand Down
1 change: 1 addition & 0 deletions torchvision/models/resnet.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@ def conv1x1(in_planes, out_planes, stride=1):

class BasicBlock(nn.Module):
expansion = 1
__constants__ = ['downsample']
Copy link
Member

@fmassa fmassa Aug 28, 2019

Choose a reason for hiding this comment

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

Can you explain a bit why this is a constant? downsample has learnable parameters.

Copy link
Contributor Author

@eellison eellison Aug 28, 2019

Choose a reason for hiding this comment

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

I can make it an empty sequential instead.

Copy link
Contributor Author

@eellison eellison Aug 28, 2019 •
edited
Loading

Choose a reason for hiding this comment

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

__constants__ only makes it a constant if it's None, otherwise it treats it as a regular submodule. It's a little bit of a confusing api, because we don't currently support Modules as first-class values and typing Optional[Module]. @driazati

Copy link
Member

@fmassa fmassa Aug 28, 2019

Choose a reason for hiding this comment

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

Ok, I think I got it. I think I prefer the __constants__ in this case, even though Optional[Module] would be nicer.


def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
base_width=64, dilation=1, norm_layer=None):
Expand Down
3 changes: 3 additions & 0 deletions torchvision/models/shufflenetv2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@


def channel_shuffle(x, groups):
# type: (torch.Tensor, int) -> torch.Tensor
batchsize, num_channels, height, width = x.data.size()
channels_per_group = num_channels // groups

Expand Down Expand Up @@ -51,6 +52,8 @@ def __init__(self, inp, oup, stride):
nn.BatchNorm2d(branch_features),
nn.ReLU(inplace=True),
)
else:
self.branch1 = nn.Sequential()
Copy link
Member

@fmassa fmassa Sep 2, 2019

Choose a reason for hiding this comment

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

I'm merging this in, but it would be great if torchscript could support Optional[nn.Module].


self.branch2 = nn.Sequential(
nn.Conv2d(inp if (self.stride > 1) else branch_features,
Expand Down