What are the allowed modification to a JIT scripted model? Modifying Dropout in jit scripted model

username: Paplham
cooked: <h3><a name="summary-question-1" class="anchor" href="#summary-question-1"></a>Summary Question</h3>
<p>In what ways can a jit-scripted model be modified? The weights can be trained. Can constants also be changed?</p>
<h3><a name="desired-use-case-2" class="anchor" href="#desired-use-case-2"></a>Desired use-case</h3>
<p>I am comparing methods for prediction uncertainty estimation. One of the commonly used baselines is Monte Carlo Dropout. To fairly evaluate the method, the dropout hyperparameters (the <code>p</code> parameter) would ideally be tuned on a hold-out set of data.</p>
<p>I am using jit-scripted models as predictors, as it allows easy incorporation of different model architectures without the need of including the model source code in the project.</p>
<h3><a name="enabling-the-dropout-of-a-scripted-model-3" class="anchor" href="#enabling-the-dropout-of-a-scripted-model-3"></a>Enabling the dropout of a scripted model</h3>
<p>Some modifications are exposed, e.g., the <code>training</code> parameter of the dropout can be set to <code>True</code> by running:</p>
<pre><code class="lang-auto">def enable_dropout(model):
    for child_name, child in model.named_children():
        if child.original_name.startswith('Dropout'):            
            child.train()
        else:
            enable_dropout(child)

model = torch.jit.load("model_gpu.pt")
enable_dropout(model)
</code></pre>
<p>However, changing the dropout probability, or replacing the dropout itself fails. See next section.</p>
<h3><a name="detailed-question-changing-the-dropout-probability-4" class="anchor" href="#detailed-question-changing-the-dropout-probability-4"></a>Detailed question: Changing the dropout probability</h3>
<p>Similarly to setting the <code>training</code> attribute, one might try to modify the <code>p</code> attribute of all dropout layers in the network by running:</p>
<pre><code class="lang-auto">def set_dropout(model, p):
    for child_name, child in model.named_children():
        if child.original_name.startswith('Dropout'): 
            setattr(child, 'p', p)
        else:
            set_dropout(child, p)

model = torch.jit.load("model_gpu.pt")
set_dropout(model, 0.5)
</code></pre>
<p>This fails with <code>RuntimeError: Can't set constant 'p' which has value: X</code>. Alternatively, one might try replacing the entire Dropout module:</p>
<pre><code class="lang-auto">def set_dropout(model, p):
    for child_name, child in model.named_children():
        if child.original_name.startswith('Dropout'): 
            setattr(model, child_name, torch.jit.script(nn.Dropout(p=p)))
        else:
            set_dropout(child, p)

model = torch.jit.load("model_gpu.pt")
set_dropout(model, 0.5)
</code></pre>
<p>This raises an error due to a mismatch in the Python compilation unit.</p>
<p>Can the desired change to the scripted model be achieved? What is the decision process for choosing which parameters are exposed to be modified and which ones are not? Is there any philosophy behind such a decision, or is it a technical constraint?</p>
<p>Thank you.</p>

username: Paplham
cooked: <p>My current workaround is to define a custom <code>Dropout</code> module as:</p>
<pre><code class="lang-auto">class _DropoutNd(nn.Module):
    __constants__ = ['inplace']
    inplace: bool

    def __init__(self, p: float = 0.5, inplace: bool = False) -&gt; None:
        super().__init__()
        if p &lt; 0 or p &gt; 1:
            raise ValueError("dropout probability has to be between 0 and 1, "
                             "but got {}".format(p))
        self.p = p
        self.inplace = inplace

    def extra_repr(self) -&gt; str:
        return 'p={}, inplace={}'.format(self.p, self.inplace)


class Dropout(_DropoutNd):
    def forward(self, input: Tensor) -&gt; Tensor:
        return F.dropout(input, self.p, self.training, self.inplace)
</code></pre>
<p>where the only modification from the original source code is that <code>p</code> is no longer defined as a constant. The approach with the <code>set_dropout</code> function then works, but I am unsure whether this modification will cause other issues down the line.</p>