Using diffetent conv2d ops with pre trained models

username: Izan_C_G
cooked: <p>Hello, I would like to know the proper practice when using different convolution operators while performing inference on pre-trained models from TorchVision, such as those found at <a href="https://pytorch.org/vision/main/models.html" rel="noopener nofollow ugc">https://pytorch.org/vision/main/models.html</a>. These models include one or more convolutional layers, but during inference, all convolutions are executed using PyTorch’s default implementation (in the case of GPU usage, via cuBLAS).</p>
<p>To change this behavior and use my custom convolution operator or one from an external library like CUTLASS (<a href="https://github.com/NVIDIA/cutlass/blob/main/examples/python/02_pytorch_extension_grouped_gemm.ipynb" rel="noopener nofollow ugc">https://github.com/NVIDIA/cutlass/blob/main/examples/python/02_pytorch_extension_grouped_gemm.ipynb</a>), I would like to know if:</p>
<ol>
<li>It is better to modify PyTorch’s source code, specifically the call to <code>torch.nn.Conv2d</code>, to import and use CUTLASS or my custom convolution operator. For this, I assume it would be necessary to recompile PyTorch or TorchVision from source.</li>
<li>It would be better to modify the model definitions in the TorchVision repository (<a href="https://github.com/pytorch/vision/tree/main/torchvision/models" rel="noopener nofollow ugc">https://github.com/pytorch/vision/tree/main/torchvision/models</a>) to use my custom convolution operator directly in their implementation.</li>
<li>The use of custom operators, as described here <a href="https://pytorch.org/tutorials/advanced/cpp_extension.html" rel="noopener nofollow ugc">https://pytorch.org/tutorials/advanced/cpp_extension.html</a>, would be a better approach.</li>
</ol>
<p>I would like to know the best practice or if there is another, more feasible or simpler approach that minimizes changes to PyTorch’s source code or is more user-friendly. Thank you.</p>

username: soulitzer
cooked: <p>I think the least invasive option is (3): register custom operators, copy paste the TorchVision model, and replace the convolutions to use your own custom operator.</p>

username: Izan_C_G
cooked: <p>Hi, <a class="mention" href="/u/soulitzer">@soulitzer</a>. Could you be more specific about how it could be done? Because of what you said, I understand that I should modify the de model internally, for instance, using the resnet definition <a href="https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py" class="inline-onebox" rel="noopener nofollow ugc">vision/torchvision/models/resnet.py at main · pytorch/vision · GitHub</a> and changing all torch.nn.con2d for my convolution. Am I right?</p>

username: soulitzer
cooked: <p>I guess what I was suggesting was more of making a literal copy of the model code from TorchVision.</p>
<p>The feasibility of this approach does depend on how self-contained the model you are targeting is.</p>

username: Izan_C_G
cooked: <p><a class="mention" href="/u/soulitzer">@soulitzer</a>, how can I make a literal copy of the model code? I mean, it is not the link I posted in the previous answer to the pytorch/vision github resnet.py file?</p>
<p>My approach is only changing the conv2d operation, nothing else. I do not intend to re-train the model, modify weights, inputs or anything else.</p>
<p>I only intend to perform a convolution with cutlass library (similarly to this <a href="https://github.com/NVIDIA/cutlass/blob/main/examples/python/03_basic_conv2d.ipynb" rel="noopener nofollow ugc">example</a> with python), but using an operator that the library allows emitting like this <a href="https://github.com/NVIDIA/cutlass/blob/main/examples/python/02_pytorch_extension_grouped_gemm.ipynb" rel="noopener nofollow ugc">gemm</a>.</p>
<p>An additional doubt is if the structure of the tensors of torch.nn.conv2d is compatible with the structure of cutlass:</p>
<pre><code class="lang-auto"># Input tensor: [N, H, W, C] under the channel-last layout
N, H, W, C = [32, 28, 28, 64]

# Weight tensor: [K, R, S, C] under the channel-last layout
K, R, S = [128, 3, 3]

# Stride, and padding
stride = (2, 2)
padding = (1, 1)
dilation = (1, 1)

# Compute the output size [N, P, Q, K]
N, P, Q, K = cutlass.Conv2d.output_size((N, H, W, C), (K, R, S, C), padding, stride, dilation)
</code></pre>
<p>Could you give some feedback about all of it?</p>

username: soulitzer
cooked: <aside class="quote no-group" data-username="Izan_C_G" data-post="5" data-topic="214367">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/izan_c_g/48/68535_2.png" class="avatar"> Izan_C_G:</div>
<blockquote>
<p>I mean, it is not the link I posted in the previous answer to the pytorch/vision github resnet.py file?</p>
</blockquote>
</aside>
<p>If you mean literally copy-pasting from the resnet.py link you provided, then that is what I mean yes.</p>

username: soulitzer
cooked: <aside class="quote no-group" data-username="Izan_C_G" data-post="5" data-topic="214367">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/izan_c_g/48/68535_2.png" class="avatar"> Izan_C_G:</div>
<blockquote>
<p>An additional doubt is if the structure of the tensors of torch.nn.conv2d is compatible with the structure of cutlass</p>
</blockquote>
</aside>
<p>I’m not sure on the specifics, but while testing you can have your nn.Conv2D replacement run each sample twice with your version and the original version to check that the results are the same.</p>