Why Conv2d is used instead of regular Patchify in ViT Pytorch?

username: omid_khalafbeigi
cooked: <p>Hi there.<br>
In Pytorch implementation of ViT, Conv2d is used over regular Patchify. in other words, researchers in <a href="https://arxiv.org/abs/2010.11929" rel="noopener nofollow ugc">An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale</a> proposed framework which receives image in a number of pieces and processes it based on self-attention mechanism. but in Pytorch version, Conv2d is used instead of that.</p>
<p>VisionTransformer Class:</p>
<pre><code class="lang-auto">def _process_input(self, x: torch.Tensor) -&gt; torch.Tensor:
        n, c, h, w = x.shape
        p = self.patch_size
        torch._assert(h == self.image_size, f"Wrong image height! Expected {self.image_size} but got {h}!")
        torch._assert(w == self.image_size, f"Wrong image width! Expected {self.image_size} but got {w}!")
        n_h = h // p
        n_w = w // p

        # (n, c, h, w) -&gt; (n, hidden_dim, n_h, n_w)
        x = self.conv_proj(x)
        # (n, hidden_dim, n_h, n_w) -&gt; (n, hidden_dim, (n_h * n_w))
        x = x.reshape(n, self.hidden_dim, n_h * n_w)

        # (n, hidden_dim, (n_h * n_w)) -&gt; (n, (n_h * n_w), hidden_dim)
        # The self attention layer expects inputs in the format (N, S, E)
        # where S is the source sequence length, N is the batch size, E is the
        # embedding dimension
        x = x.permute(0, 2, 1)

        return x
</code></pre>
<pre><code class="lang-auto">def forward(self, x: torch.Tensor):
        # Reshape and permute the input tensor
        x = self._process_input(x)
        n = x.shape[0]

        # Expand the class token to the full batch
        batch_class_token = self.class_token.expand(n, -1, -1)
        x = torch.cat([batch_class_token, x], dim=1)

        x = self.encoder(x)

        # Classifier "token" as used by standard language architectures
        x = x[:, 0]

        x = self.heads(x)

        return x
</code></pre>
<p>Can anyone explain for me why Conv2d is used? why not regular Patchify?!<br>
Thanks.</p>

username: J_Johnson
cooked: <pre><code class="lang-auto">if conv_stem_configs is not None:
         &gt;&gt;&gt;# As per https://arxiv.org/abs/2106.14881&lt;&lt;&lt;
            seq_proj = nn.Sequential()
            prev_channels = 3
            for i, conv_stem_layer_config in enumerate(conv_stem_configs):
                seq_proj.add_module(
                    f"conv_bn_relu_{i}",
                    Conv2dNormActivation(
                        in_channels=prev_channels,
                        out_channels=conv_stem_layer_config.out_channels,
                        kernel_size=conv_stem_layer_config.kernel_size,
                        stride=conv_stem_layer_config.stride,
                        norm_layer=conv_stem_layer_config.norm_layer,
                        activation_layer=conv_stem_layer_config.activation_layer,
                    ),
                )
                prev_channels = conv_stem_layer_config.out_channels
            seq_proj.add_module(
                "conv_last", nn.Conv2d(in_channels=prev_channels, out_channels=hidden_dim, kernel_size=1)
            )
            self.conv_proj: nn.Module = seq_proj
        else:
            self.conv_proj = nn.Conv2d(
                in_channels=3, out_channels=hidden_dim, kernel_size=patch_size, stride=patch_size
            )
</code></pre>
<p>It says in the comments to see this paper: <a href="https://arxiv.org/abs/2106.14881" class="inline-onebox" rel="noopener nofollow ugc">[2106.14881] Early Convolutions Help Transformers See Better</a></p>
<p>I’m sure they did it because the data shows it to be an improvement over not doing so.</p>
<p>Seems the original Google ViT model had a “patchify stem”, which had a conv layer with a large 16x16 kernel.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/f/e/fe6e66ea5c293d86f4ed0425c946311a0860d1a9.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/fe6e66ea5c293d86f4ed0425c946311a0860d1a9" title="Screenshot_2023-03-02-16-26-45-07_c0fc23361849dfc0a4b4c6dd2adcac79"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/f/e/fe6e66ea5c293d86f4ed0425c946311a0860d1a9_2_690x310.jpeg" alt="Screenshot_2023-03-02-16-26-45-07_c0fc23361849dfc0a4b4c6dd2adcac79" data-base62-sha1="AiNOo5grdVghf5Giw6Kw9Fp0vep" width="690" height="310" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/f/e/fe6e66ea5c293d86f4ed0425c946311a0860d1a9_2_690x310.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/f/e/fe6e66ea5c293d86f4ed0425c946311a0860d1a9_2_1035x465.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/f/e/fe6e66ea5c293d86f4ed0425c946311a0860d1a9_2_1380x620.jpeg 2x" data-dominant-color="CCCBCB"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot_2023-03-02-16-26-45-07_c0fc23361849dfc0a4b4c6dd2adcac79</span><span class="informations">2400×1080 188 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>But specifying this option in the Pytorch ViT changes that stem kernel to be a custom size, which the above paper demonstrates that size 3x3 followed by a 1x1 performs better.</p>

username: J_Johnson
cooked: <p>Thought it might be helpful to add how to build a custom patchify stem, given the docs do not contain an example. This repeats what they did in the paper:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/5/5/5548452c8b7866df6155162123d7a9e1bc9e11d9.png" data-download-href="https://discuss.pytorch.org/uploads/default/5548452c8b7866df6155162123d7a9e1bc9e11d9" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/5/5/5548452c8b7866df6155162123d7a9e1bc9e11d9_2_690x223.png" alt="image" data-base62-sha1="carqeMEvWA7SaRUjKzjKxDbM6Kd" width="690" height="223" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/5/5/5548452c8b7866df6155162123d7a9e1bc9e11d9_2_690x223.png, https://discuss.pytorch.org/uploads/default/optimized/3X/5/5/5548452c8b7866df6155162123d7a9e1bc9e11d9_2_1035x334.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/5/5/5548452c8b7866df6155162123d7a9e1bc9e11d9_2_1380x446.png 2x" data-dominant-color="E9E8E6"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1513×490 170 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<pre><code class="lang-auto">from torchvision.models.vision_transformer import ConvStemConfig, vit_b_16

n = 4
stemconfig = [ConvStemConfig(out_channels = 64, kernel_size = 3 , stride = 2) for i in range(n)]

model = vit_b_16(pretrained=False, conv_stem_configs=stemconfig)
</code></pre>

username: omid_khalafbeigi
cooked: <p>Thank you so much for this code example!<br>
Do you have any idea why researchers not mentioned using convolution instead of regular patchify in ViT original paper? I mean according to below figure and its description, regular patchify is used…<br>
<strong>Maybe their mean about linearly embedding is convolution operation.</strong></p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/a/f/af7bc934d9b9dc58e47fe0bb308720cfc2717966.png" data-download-href="https://discuss.pytorch.org/uploads/default/af7bc934d9b9dc58e47fe0bb308720cfc2717966" title="Screenshot from 2023-03-03 15-12-13"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/a/f/af7bc934d9b9dc58e47fe0bb308720cfc2717966_2_690x386.png" alt="Screenshot from 2023-03-03 15-12-13" data-base62-sha1="p2oLAo0y9eQ2EXgqvfjdYzCcv3g" width="690" height="386" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/a/f/af7bc934d9b9dc58e47fe0bb308720cfc2717966_2_690x386.png, https://discuss.pytorch.org/uploads/default/optimized/3X/a/f/af7bc934d9b9dc58e47fe0bb308720cfc2717966_2_1035x579.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/a/f/af7bc934d9b9dc58e47fe0bb308720cfc2717966.png 2x" data-dominant-color="E6E4E3"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot from 2023-03-03 15-12-13</span><span class="informations">1257×704 148 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

username: J_Johnson
cooked: <p>If you flatten each 16x16 segment and apply a linear layer to turn each into a size of <code>hidden_dim</code>, you’re effectively doing the same as applying a Conv2d with kernel and stride of 16x16 with output of <code>hidden_dim</code>. But a convolution is likely a preferred operation, if I recall correctly, because it’s faster for GPUs.</p>
<p>However, I cannot speak to the actual reasoning of the developers and the above is my own speculation regarding their reasoning.</p>

username: hubert10
cooked: <p>Hi Omid, as far as the patchification vs Convolutions of/over the input image is concerned, there are quite few reasons why one would choose convolutions over the standard patchification:</p>
<ol>
<li>with convolutions, the local spatial information is retained via convolution kernels</li>
<li>for low-resolution images, patching generated smaller patches, which can lead to a loss of fine-grained features.</li>
</ol>