Set Dropout to be True for all the different layers of a Transformer model!

Mathedelic: <p>I am working on image classification using a Transformer! The problem is that of overfitting, I am getting training accuracy of 1, but validation and test accuracy is around 70%. I wanted to ask how can I set dropout to be true in different layers of the model. Here is the architecture of the model :<br>
VisionTransformer(<br>
(patch_embed): PatchEmbed(<br>
(proj): Conv2d(3, 768, kernel_size=(16, 16), stride=(16, 16))<br>
(norm): Identity()<br>
)<br>
(pos_drop): Dropout(p=0.5, inplace=True)<br>
(blocks): Sequential(<br>
(0): Block(<br>
(norm1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)<br>
(attn): Attention(<br>
(qkv): Linear(in_features=768, out_features=2304, bias=True)<br>
(attn_drop): Dropout(p=0.0, inplace=False)<br>
(proj): Linear(in_features=768, out_features=768, bias=True)<br>
(proj_drop): Dropout(p=0.0, inplace=False)<br>
)<br>
(drop_path): Identity()<br>
(norm2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)<br>
(mlp): Mlp(<br>
(fc1): Linear(in_features=768, out_features=3072, bias=True)<br>
(act): GELU()<br>
(fc2): Linear(in_features=3072, out_features=768, bias=True)<br>
(drop): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
(1): Block(<br>
(norm1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)<br>
(attn): Attention(<br>
(qkv): Linear(in_features=768, out_features=2304, bias=True)<br>
(attn_drop): Dropout(p=0.0, inplace=False)<br>
(proj): Linear(in_features=768, out_features=768, bias=True)<br>
(proj_drop): Dropout(p=0.0, inplace=False)<br>
)<br>
(drop_path): Identity()<br>
(norm2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)<br>
(mlp): Mlp(<br>
(fc1): Linear(in_features=768, out_features=3072, bias=True)<br>
(act): GELU()<br>
(fc2): Linear(in_features=3072, out_features=768, bias=True)<br>
(drop): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
(2): Block(<br>
(norm1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)<br>
(attn): Attention(<br>
(qkv): Linear(in_features=768, out_features=2304, bias=True)<br>
(attn_drop): Dropout(p=0.0, inplace=False)<br>
(proj): Linear(in_features=768, out_features=768, bias=True)<br>
(proj_drop): Dropout(p=0.0, inplace=False)<br>
)<br>
(drop_path): Identity()<br>
(norm2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)<br>
(mlp): Mlp(<br>
(fc1): Linear(in_features=768, out_features=3072, bias=True)<br>
(act): GELU()<br>
(fc2): Linear(in_features=3072, out_features=768, bias=True)<br>
(drop): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
(3): Block(<br>
(norm1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)<br>
(attn): Attention(<br>
(qkv): Linear(in_features=768, out_features=2304, bias=True)<br>
(attn_drop): Dropout(p=0.0, inplace=False)<br>
(proj): Linear(in_features=768, out_features=768, bias=True)<br>
(proj_drop): Dropout(p=0.0, inplace=False)<br>
)<br>
(drop_path): Identity()<br>
(norm2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)<br>
(mlp): Mlp(<br>
(fc1): Linear(in_features=768, out_features=3072, bias=True)<br>
(act): GELU()<br>
(fc2): Linear(in_features=3072, out_features=768, bias=True)<br>
(drop): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
(4): Block(<br>
(norm1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)<br>
(attn): Attention(<br>
(qkv): Linear(in_features=768, out_features=2304, bias=True)<br>
(attn_drop): Dropout(p=0.0, inplace=False)<br>
(proj): Linear(in_features=768, out_features=768, bias=True)<br>
(proj_drop): Dropout(p=0.0, inplace=False)<br>
)<br>
(drop_path): Identity()<br>
(norm2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)<br>
(mlp): Mlp(<br>
(fc1): Linear(in_features=768, out_features=3072, bias=True)<br>
(act): GELU()<br>
(fc2): Linear(in_features=3072, out_features=768, bias=True)<br>
(drop): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
(5): Block(<br>
(norm1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)<br>
(attn): Attention(<br>
(qkv): Linear(in_features=768, out_features=2304, bias=True)<br>
(attn_drop): Dropout(p=0.0, inplace=False)<br>
(proj): Linear(in_features=768, out_features=768, bias=True)<br>
(proj_drop): Dropout(p=0.0, inplace=False)<br>
)<br>
(drop_path): Identity()<br>
(norm2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)<br>
(mlp): Mlp(<br>
(fc1): Linear(in_features=768, out_features=3072, bias=True)<br>
(act): GELU()<br>
(fc2): Linear(in_features=3072, out_features=768, bias=True)<br>
(drop): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
(6): Block(<br>
(norm1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)<br>
(attn): Attention(<br>
(qkv): Linear(in_features=768, out_features=2304, bias=True)<br>
(attn_drop): Dropout(p=0.0, inplace=False)<br>
(proj): Linear(in_features=768, out_features=768, bias=True)<br>
(proj_drop): Dropout(p=0.0, inplace=False)<br>
)<br>
(drop_path): Identity()<br>
(norm2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)<br>
(mlp): Mlp(<br>
(fc1): Linear(in_features=768, out_features=3072, bias=True)<br>
(act): GELU()<br>
(fc2): Linear(in_features=3072, out_features=768, bias=True)<br>
(drop): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
(7): Block(<br>
(norm1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)<br>
(attn): Attention(<br>
(qkv): Linear(in_features=768, out_features=2304, bias=True)<br>
(attn_drop): Dropout(p=0.0, inplace=False)<br>
(proj): Linear(in_features=768, out_features=768, bias=True)<br>
(proj_drop): Dropout(p=0.0, inplace=False)<br>
)<br>
(drop_path): Identity()<br>
(norm2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)<br>
(mlp): Mlp(<br>
(fc1): Linear(in_features=768, out_features=3072, bias=True)<br>
(act): GELU()<br>
(fc2): Linear(in_features=3072, out_features=768, bias=True)<br>
(drop): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
(8): Block(<br>
(norm1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)<br>
(attn): Attention(<br>
(qkv): Linear(in_features=768, out_features=2304, bias=True)<br>
(attn_drop): Dropout(p=0.0, inplace=False)<br>
(proj): Linear(in_features=768, out_features=768, bias=True)<br>
(proj_drop): Dropout(p=0.0, inplace=False)<br>
)<br>
(drop_path): Identity()<br>
(norm2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)<br>
(mlp): Mlp(<br>
(fc1): Linear(in_features=768, out_features=3072, bias=True)<br>
(act): GELU()<br>
(fc2): Linear(in_features=3072, out_features=768, bias=True)<br>
(drop): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
(9): Block(<br>
(norm1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)<br>
(attn): Attention(<br>
(qkv): Linear(in_features=768, out_features=2304, bias=True)<br>
(attn_drop): Dropout(p=0.0, inplace=False)<br>
(proj): Linear(in_features=768, out_features=768, bias=True)<br>
(proj_drop): Dropout(p=0.0, inplace=False)<br>
)<br>
(drop_path): Identity()<br>
(norm2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)<br>
(mlp): Mlp(<br>
(fc1): Linear(in_features=768, out_features=3072, bias=True)<br>
(act): GELU()<br>
(fc2): Linear(in_features=3072, out_features=768, bias=True)<br>
(drop): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
(10): Block(<br>
(norm1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)<br>
(attn): Attention(<br>
(qkv): Linear(in_features=768, out_features=2304, bias=True)<br>
(attn_drop): Dropout(p=0.0, inplace=False)<br>
(proj): Linear(in_features=768, out_features=768, bias=True)<br>
(proj_drop): Dropout(p=0.0, inplace=False)<br>
)<br>
(drop_path): Identity()<br>
(norm2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)<br>
(mlp): Mlp(<br>
(fc1): Linear(in_features=768, out_features=3072, bias=True)<br>
(act): GELU()<br>
(fc2): Linear(in_features=3072, out_features=768, bias=True)<br>
(drop): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
(11): Block(<br>
(norm1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)<br>
(attn): Attention(<br>
(qkv): Linear(in_features=768, out_features=2304, bias=True)<br>
(attn_drop): Dropout(p=0.0, inplace=False)<br>
(proj): Linear(in_features=768, out_features=768, bias=True)<br>
(proj_drop): Dropout(p=0.0, inplace=False)<br>
)<br>
(drop_path): Identity()<br>
(norm2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)<br>
(mlp): Mlp(<br>
(fc1): Linear(in_features=768, out_features=3072, bias=True)<br>
(act): GELU()<br>
(fc2): Linear(in_features=3072, out_features=768, bias=True)<br>
(drop): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
)<br>
(norm): LayerNorm((768,), eps=1e-06, elementwise_affine=True)<br>
(pre_logits): Identity()<br>
(head): Linear(in_features=768, out_features=2, bias=True)</p>
<p>I want to set dropout to be 0.5 in all the different layers. Starting from the first layer when I do : <code>model_ft._modules["pos_drop"] = nn.Dropout(0.5, inplace=True)</code>, it works for the first instance of the dropout but when I want to do the same thing for the second dropout and I try <code>model_ft._modules["blocks"].attn.proj_drop = nn.Dropout(0.5, inplace=True)</code>, it throws an error.</p>
<p>The real problem is that I don’t know how to access the dropout layers in the network and set them all to non-zero values.</p>
<p>I would be very thankful to you if you could help me with how to access the different layers of  the model and set dropout to be true in all of them.</p>

ptrblck: <p>You could use <code>model.apply</code> to iterate all modules, check if the current module is an <code>nn.Dropout</code>, and change the <code>p</code> value as seen here:</p>
<pre><code class="lang-python">class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(10, 10)
        self.drop1 = nn.Dropout(p=0.0)
        self.fc2 = nn.Linear(10, 10)
        self.drop2 = nn.Dropout(p=0.0)
        self.fc3 = nn.Linear(10, 10)
        self.drop3 = nn.Dropout(p=0.0)
        
model = MyModel()
print(model)

def set_dropout_p(m, p):
    if isinstance(m, nn.Dropout):
        m.p = p

model.apply(lambda m: set_dropout_p(m, p=0.5))
print(model)
</code></pre>

Mathedelic: <p>YOU ARE A GOD ! Thank you so much it worked!<br>
<img src="https://discuss.pytorch.org/images/emoji/apple/om.png?v=10" title=":om:" class="emoji only-emoji" alt=":om:"></p>

ankita.buntolia: <p>How can we make sure that <strong>dropouts</strong> is only used in training and not in validation? How can we set in the network only for training?</p>
<p>This is how I did it based on dropout documentation!</p>
<pre><code class="lang-auto">#####  U-Net Model
## 
class DoubleConv(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(DoubleConv, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, 3, 1, 1, bias=False),
            nn.Dropout2d(p=0.5),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, 3, 1, 1, bias=False),
            nn.Dropout2d(p=0.5),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
        )

    def forward(self, x):
        return self.conv(x)

class UNET(nn.Module):
    def __init__(
            self, in_channels=1, out_channels=3, features=[64, 128, 256, 512],
    ):
        super(UNET, self).__init__()
        self.ups = nn.ModuleList()
        self.downs = nn.ModuleList()
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)

        # Down part of UNET
        for feature in features:
            self.downs.append(DoubleConv(in_channels, feature))
            in_channels = feature

        # Up part of UNET
        for feature in reversed(features):
            self.ups.append(
                nn.ConvTranspose2d(
                    feature*2, feature, kernel_size=2, stride=2,
                )
            )
            self.ups.append(DoubleConv(feature*2, feature))

        self.bottleneck = DoubleConv(features[-1], features[-1]*2)
        self.final_conv = nn.Conv2d(features[0], out_channels, kernel_size=1)
        # softmax
        self.softmax=nn.Softmax(dim=1)

    def forward(self, x):
      skip_connections = []

      for down in self.downs:
        x = down(x)
        skip_connections.append(x)
        x = self.pool(x)

      x = self.bottleneck(x)
      skip_connections = skip_connections[::-1]

      for idx in range(0, len(self.ups), 2):
        x = self.ups[idx](x)
        skip_connection = skip_connections[idx//2]

        if x.shape != skip_connection.shape:
          x = TF.resize(x, size=skip_connection.shape[2:])

        concat_skip = torch.cat((skip_connection, x), dim=1)
        x = self.ups[idx+1](concat_skip)
        # x = self.final_conv(x)
        # x = self.softmax(x)

      # return self.final_conv(x)


      # return self.softmax(self.final_conv(x))
      return self.final_conv(x)
</code></pre>

ptrblck: <p><code>nn.Dropout</code> layer will automatically switch their behavior based on the <code>self.training</code> flag, which is used internally in their <code>forward</code> method.<br>
Calling <code>model.train()</code> and <code>model.eval()</code> will switch this attribute to <code>True</code> or <code>False</code>, respectively.</p>

ankita.buntolia: <p>Thank you very much!</p>