make shufflenet and resnet scriptable #1270
Conversation
| @@ -34,6 +34,7 @@ def conv1x1(in_planes, out_planes, stride=1): | |||
|
|
|||
| class BasicBlock(nn.Module): | |||
| expansion = 1 | |||
| __constants__ = ['downsample'] | |||
There was a problem hiding this comment.
Can you explain a bit why this is a constant? downsample has learnable parameters.
There was a problem hiding this comment.
I can make it an empty sequential instead.
There was a problem hiding this comment.
__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
There was a problem hiding this comment.
Ok, I think I got it. I think I prefer the __constants__ in this case, even though Optional[Module] would be nicer.
eellison
changed the title
|
CI is failing, probably because we might need to use PyTorch nightly now? Can you try changing Line 29 in 0bd7080 |
|
I've merged another PR that changed the CI to use pytorch nightly, so let me re-run the CI |
|
@fmassa thanks! |
|
@eellison there is a linter error due to the type annotation not finding |
yep will do |
Codecov Report
@@ Coverage Diff @@
## master #1270 +/- ##
==========================================
- Coverage 65.56% 65.48% -0.08%
==========================================
Files 74 74
Lines 5784 5786 +2
Branches 884 884
==========================================
- Hits 3792 3789 -3
- Misses 1726 1730 +4
- Partials 266 267 +1
Continue to review full report at Codecov.
|
| @@ -51,6 +52,8 @@ def __init__(self, inp, oup, stride): | |||
| nn.BatchNorm2d(branch_features), | |||
| nn.ReLU(inplace=True), | |||
| ) | |||
| else: | |||
| self.branch1 = nn.Sequential() | |||
There was a problem hiding this comment.
I'm merging this in, but it would be great if torchscript could support Optional[nn.Module].
ShuffleNet:
branch1was only conditionally definedResnet:
downsampleto constants so that script can compile when it is None