Optional submodule in ScriptedModule

username: apisarek
cooked: <p>Hi! How should I declare an optional BatchNorm submodule in a valid TorchScript way?</p>
<pre><code class="lang-auto">class Downsample(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: int,
        apply_batchnorm: bool = False,
    ):
        super(Downsample, self).__init__()
        self.conv = nn.Conv2d(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=kernel_size,
            stride=2,
            bias=False,
            padding=1,
        )
        self.apply_batchnorm = apply_batchnorm
        if self.apply_batchnorm:
            self.bn = nn.BatchNorm2d(out_channels)

    def forward(self, x):
        x = self.conv(x)
        if self.apply_batchnorm:
            x = self.bn(x)
        x = F.leaky_relu(x, 0.3)
        return x
</code></pre>
<p>When trying to run <code>torch.jit.script(Downsample(1, 2, 3))</code> I get following error.</p>
<pre><code class="lang-auto">RuntimeError: 
Module 'Downsample' has no attribute 'bn' :
  File "path/to/file.py", line 100
        x = self.conv(x)
        if self.apply_batchnorm:
            x = self.bn(x)
                ~~~~~~~ &lt;--- HERE
        x = F.leaky_relu(x, 0.3)
        return x
</code></pre>
<p>I understand that in TorchScript every variable must have single static type. Adding <code>bn: Optional[nn.BatchNorm2d]</code> to the class definition does not help.</p>

username: driazati
cooked: <p>This is close, the issue is that when <code>self.apply_batchnorm</code> is false, there is no <code>bn</code> attribute on the module, so it cannot be accessed / checked. So</p>
<pre><code class="lang-python">        if self.apply_batchnorm:
            self.bn = nn.BatchNorm2d(out_channels)
</code></pre>
<p>turns into</p>
<pre><code class="lang-python">        if self.apply_batchnorm:
            self.bn = nn.BatchNorm2d(out_channels)
        else:
            self.bn = None
</code></pre>
<p>Another piece is that to call an optional module, the compiler must be able to figure out that it is not <code>None</code> when it is called, so instead of <code>if self.apply_batchnorm</code> you have to do <code>if self.bn is not None</code>.</p>
<p>This is the full working example</p>
<pre><code class="lang-python">class Downsample(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: int,
        apply_batchnorm: bool = False,
    ):
        super(Downsample, self).__init__()
        self.conv = nn.Conv2d(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=kernel_size,
            stride=2,
            bias=False,
            padding=1,
        )
        self.apply_batchnorm = apply_batchnorm
        if self.apply_batchnorm:
            self.bn = nn.BatchNorm2d(out_channels)
        else:
            self.bn = None

    def forward(self, x):
        x = self.conv(x)
        if self.bn is not None:
            x = self.bn(x)
        x = F.leaky_relu(x, 0.3)
        return x


torch.jit.script(Downsample(1, 2, 3))
</code></pre>

username: G.M
cooked: <p>U can also use <code>nn.Identity</code> to replace the batchnorm. For example, in <code>__init__</code>, u can write:</p>
<pre><code class="lang-auto">if self.apply_batchnorm:
    self.bn = nn.BatchNorm2d(out_channels)
else:
    self.bn = nn.Identity()
</code></pre>

username: apisarek
cooked: <p>Thank you driazati and G.M, I have also found another solution in Github issues in the meantime.</p>
<aside class="onebox githubpullrequest">
  <header class="source">
      <a href="https://github.com/pytorch/vision/pull/1270/files/0f8a35d12d7b17776b2feb561f73e91fcedc8026#diff-bab13aad3b16866a3936568fb6fcbe70" target="_blank" rel="nofollow noopener">github.com/pytorch/vision</a>
  </header>
  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Pull Request">
    <svg width="60" height="60" class="github-icon" viewbox="0 0 12 16" aria-hidden="true"><path d="M11 11.28V5c-.03-.78-.34-1.47-.94-2.06C9.46 2.35 8.78 2.03 8 2H7V0L4 3l3 3V4h1c.27.02.48.11.69.31.21.2.3.42.31.69v6.28A1.993 1.993 0 0 0 10 15a1.993 1.993 0 0 0 1-3.72zm-1 2.92c-.66 0-1.2-.55-1.2-1.2 0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2 0 .65-.55 1.2-1.2 1.2zM4 3c0-1.11-.89-2-2-2a1.993 1.993 0 0 0-1 3.72v6.56A1.993 1.993 0 0 0 2 15a1.993 1.993 0 0 0 1-3.72V4.72c.59-.34 1-.98 1-1.72zm-.8 10c0 .66-.55 1.2-1.2 1.2-.65 0-1.2-.55-1.2-1.2 0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2zM2 4.2C1.34 4.2.8 3.65.8 3c0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2 0 .65-.55 1.2-1.2 1.2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/pytorch/vision/pull/1270" target="_blank" rel="nofollow noopener">make shufflenet and resnet scriptable</a>
    </h4>

    <div class="branches">
      <code>pytorch:master</code> ← <code>eellison:shufflenet</code>
    </div>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2019-08-28" data-time="16:11:54" data-timezone="UTC">04:11PM - 28 Aug 19 UTC</span>
      </div>

      <div class="user">
        <a href="https://github.com/eellison" target="_blank" rel="nofollow noopener">
          <img alt="eellison" src="https://avatars2.githubusercontent.com/u/11477974?v=4" class="onebox-avatar-inline" width="20" height="20">
          eellison
        </a>
      </div>

      <div class="lines" title="6 commits changed 3 files with 6 additions and 2 deletions">
        <a href="https://github.com/pytorch/vision/pull/1270/files" target="_blank" rel="nofollow noopener">
          <span class="added">+6</span>
          <span class="removed">-2</span>
        </a>
      </div>
    </div>

  </div>
</div>
  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>

<aside class="onebox githubissue">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/issues/26067" target="_blank" rel="nofollow noopener">github.com/pytorch/pytorch</a>
  </header>
  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Issue">
	  <svg width="60" height="60" class="github-icon" viewbox="0 0 14 16" aria-hidden="true"><path d="M7 2.3c3.14 0 5.7 2.56 5.7 5.7s-2.56 5.7-5.7 5.7A5.71 5.71 0 0 1 1.3 8c0-3.14 2.56-5.7 5.7-5.7zM7 1C3.14 1 0 4.14 0 8s3.14 7 7 7 7-3.14 7-7-3.14-7-7-7zm1 3H6v5h2V4zm0 6H6v2h2v-2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/pytorch/pytorch/issues/26067" target="_blank" rel="nofollow noopener">TorchScript doesn't handle submodule being None</a>
    </h4>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2019-09-11" data-time="23:37:20" data-timezone="UTC">11:37PM - 11 Sep 19 UTC</span>
      </div>

        <div class="date">
          closed <span class="discourse-local-date" data-format="ll" data-date="2020-01-27" data-time="07:05:52" data-timezone="UTC">07:05AM - 27 Jan 20 UTC</span>
        </div>

      <div class="user">
        <a href="https://github.com/dzhulgakov" target="_blank" rel="nofollow noopener">
          <img alt="dzhulgakov" src="https://avatars2.githubusercontent.com/u/17890620?v=4" class="onebox-avatar-inline" width="20" height="20">
          dzhulgakov
        </a>
      </div>
    </div>
  </div>
</div>

<div class="github-row">
  <p class="github-content">🐛 Bug
If one of the submodules is None, TorchScript doesn't recognize it as an attribute, unless explicitly specified in __constants__ lists.
Obviously,...</p>
</div>

<div class="labels">
    <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">jit</span>
    <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">jit-backlog</span>
</div>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>