How to reduce latency in ConvTranspose2d?

username: hankyul
cooked: <p><strong>Question</strong>: How to reduce latency(computational cost) in <code>ConvTranspose2d</code>?</p>
<p>I compared latency between <code>nn.Conv2d</code>, <code>Upsample(mode=nearest)</code>, <code>ConvTranspose2d</code>, <code>Upsample(mode=bilinear)</code> in different batch size. It looks like deconvolution operation takes a lot of time in large batch (=64). Is there a way to reduce computation time for deconvolution?</p>
<p>Thank you</p>
<p><strong>Comparison Table</strong></p>
<pre><code class="lang-bash">[------------------------------------- upsample module comparison ------------------------------------]
                        |  nearest(scale=2)  |  conv(kernel=3)  |  deconv(scale=2)  |  bilinear(scale=2)
4 threads: --------------------------------------------------------------------------------------------
      [8, 2048, 8, 8]   |        30.4        |       94.6      |        134.6      |        4838.2     
      [16, 2048, 8, 8]  |        58.4        |       94.3      |        249.4      |        9682.3     
      [32, 2048, 8, 8]  |       111.6        |      119.1      |        480.1      |       19434.2     
      [64, 2048, 8, 8]  |       219.9        |      206.1      |      12399.6      |       38844.8     

Times are in microseconds (us).
</code></pre>
<p><strong>Comparison Code</strong></p>
<pre><code class="lang-python">import os

os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
os.environ['CUDA_VISIBLE_DEVICES'] = '4'

import torch
import torch.nn as nn
import torch.utils.benchmark as benchmark
import torchvision.models as models
from itertools import product


def get_upsample_module(mode, upsample, ch):
    if mode == 'deconv':
        return nn.Sequential(
            torch.nn.ConvTranspose2d(ch, ch, upsample, stride=upsample, dilation=1, groups=ch, bias=False),
            torch.nn.BatchNorm2d(ch)
        ).cuda()
    else:
        return nn.Upsample(scale_factor=upsample, mode=mode).cuda()

def get_downsample_module(kernel, ch):
    return nn.Sequential(
        torch.nn.Conv2d(ch, ch, kernel, stride=1, padding=1, dilation=1, groups=ch, bias=False),
        torch.nn.BatchNorm2d(ch)
    ).cuda()


results = []

batch_size = [8, 16, 32, 64]
channel_size = [2048]
image_size = [8]
mode = ['nearest', 'conv', 'deconv', 'bilinear']
scale_factors = [2]

for b, c, n in product(batch_size, channel_size, image_size):
    label = 'upsample module comparison'
    sub_label = f'[{b}, {c}, {n}, {n}]'
    x = torch.rand((b, c, n, n)).cuda()

    for method, upsample in product(mode, scale_factors):
        if method == 'conv':
            upsample += 1
            model = get_downsample_module(upsample, c)
        else:
            model = get_upsample_module(method, upsample, c)
        with torch.cuda.amp.autocast():
            results.append(benchmark.Timer(
                stmt='model(x)',
                setup='from __main__ import model',
                globals={'x': x},
                label=label,
                sub_label=sub_label,
                description=f"{method}(scale={upsample})",
                num_threads=4
            ).blocked_autorange(min_run_time=1))

compare = benchmark.Compare(results)
compare.colorize()
compare.print()
</code></pre>
<p><strong>My Environment</strong>:</p>
<ul>
<li>OS: Linux 18.04</li>
<li>Python: 3.7.11</li>
<li>Pytorch: 1.10.0</li>
<li>CUDA: 11.3</li>
</ul>

username: ptrblck
cooked: <p>You could add <code>torch.backends.cudnn.benchmark = True</code> to let cuDNN profile different kernels and select the fastest one and use <code>to(memory_format=torch.channels_last)</code> since you are using <code>amp</code>, which could give another speedup if permutations are avoided.</p>

username: hankyul
cooked: <p>Thank you <a class="mention" href="/u/ptrblck">@ptrblck</a> I solve latency problem by following your suggestion!!!</p>
<p>For those having simmiar problem, I recommend to use <code>to(memory_format=torch.channels_last)</code></p>
<h3>
<a name="benchmark-amp-memory_format-experiment-results-1" class="anchor" href="#benchmark-amp-memory_format-experiment-results-1"></a>Benchmark, AMP, memory_format experiment results</h3>
<p><code>nn.Conv2d</code> , <code>Upsample(mode=nearest)</code> , <code>ConvTranspose2d</code> , <code>Upsample(mode=bilinear)</code> latency results on different settings.</p>
<ol>
<li>
<p>Default (bechmark=False, AMP=False, memory_format=contiguos)</p>
<pre><code class="lang-bash">[-------------------------------------- upsample module comparison -------------------------------------]
                          |  nearest(scale=2)  |  conv(scale=3)  |  deconv(scale=2)  |  bilinear(scale=2)
4 threads: ----------------------------------------------------------------------------------------------
      [8, 2048, 16, 16]   |       110.4        |       144.6     |        492.8      |        4891.5     
      [16, 2048, 16, 16]  |       216.6        |       276.8     |       1128.5      |        9821.6     
      [32, 2048, 16, 16]  |       433.9        |       541.6     |       2255.2      |       19672.0     
      [64, 2048, 16, 16]  |       864.3        |      1214.2     |       4535.2      |       39559.5     

Times are in microseconds (us).
</code></pre>
</li>
<li>
<p>Nothing changed (Benchmark=True, AMP=False, memory_format=contiguous)</p>
<pre><code class="lang-bash">[-------------------------------------- upsample module comparison -------------------------------------]
                          |  nearest(scale=2)  |  conv(scale=3)  |  deconv(scale=2)  |  bilinear(scale=2)
4 threads: ----------------------------------------------------------------------------------------------
      [8, 2048, 16, 16]   |       110.5        |       145.4     |        494.8      |        4906.8     
      [16, 2048, 16, 16]  |       216.9        |       277.9     |       1131.3      |        9825.4     
      [32, 2048, 16, 16]  |       434.4        |       543.0     |       2266.0      |       19766.6     
      [64, 2048, 16, 16]  |       865.5        |      1211.6     |       4539.3      |       39544.8     

Times are in microseconds (us).
</code></pre>
</li>
<li>
<p><strong>Conv (<img src="https://discuss.pytorch.org/images/emoji/apple/small_red_triangle_down.png?v=12" title=":small_red_triangle_down:" class="emoji" alt=":small_red_triangle_down:">)</strong> <strong>Deconv (<img src="https://discuss.pytorch.org/images/emoji/apple/small_red_triangle.png?v=12" title=":small_red_triangle:" class="emoji" alt=":small_red_triangle:">)</strong> (Benchmark=True, AMP=True, memory_format=contiguous)</p>
<pre><code class="lang-bash">[-------------------------------------- upsample module comparison -------------------------------------]
                          |  nearest(scale=2)  |  conv(scale=3)  |  deconv(scale=2)  |  bilinear(scale=2)
4 threads: ----------------------------------------------------------------------------------------------
      [8, 2048, 16, 16]   |       110.5        |      162.3      |        443.0      |        4907.9     
      [16, 2048, 16, 16]  |       217.0        |      169.6      |      11440.2      |        9828.0     
      [32, 2048, 16, 16]  |       434.5        |      320.2      |      11440.9      |       19671.5     
      [64, 2048, 16, 16]  |       865.7        |      619.6      |      14415.1      |       39541.3     

Times are in microseconds (us).
</code></pre>
</li>
<li>
<p><strong>Nearest</strong> <strong>(<img src="https://discuss.pytorch.org/images/emoji/apple/small_red_triangle.png?v=12" title=":small_red_triangle:" class="emoji" alt=":small_red_triangle:">)</strong> <strong>Conv</strong> <strong>(<img src="https://discuss.pytorch.org/images/emoji/apple/small_red_triangle_down.png?v=12" title=":small_red_triangle_down:" class="emoji" alt=":small_red_triangle_down:">)</strong> <strong>Deconv</strong> <strong>(<img src="https://discuss.pytorch.org/images/emoji/apple/small_red_triangle.png?v=12" title=":small_red_triangle:" class="emoji" alt=":small_red_triangle:">)</strong> <strong>Bilinear</strong> <strong>(<img src="https://discuss.pytorch.org/images/emoji/apple/small_red_triangle_down.png?v=12" title=":small_red_triangle_down:" class="emoji" alt=":small_red_triangle_down:">)</strong> (Benchmark=True, AMP=True, memory_format=channels_last)</p>
<pre><code class="lang-bash">[-------------------------------------- upsample module comparison -------------------------------------]
                          |  nearest(scale=2)  |  conv(scale=3)  |  deconv(scale=2)  |  bilinear(scale=2)
4 threads: ----------------------------------------------------------------------------------------------
      [8, 2048, 16, 16]   |        247.2       |      154.8      |        825.5      |         336.3     
      [16, 2048, 16, 16]  |        492.0       |      170.3      |       1677.0      |         675.8     
      [32, 2048, 16, 16]  |        979.6       |      327.0      |       3498.6      |        1347.3     
      [64, 2048, 16, 16]  |       1971.4       |      675.1      |       6861.0      |        2703.2     

Times are in microseconds (us).
</code></pre>
</li>
<li>
<p><strong>Nearest</strong> <strong>(<img src="https://discuss.pytorch.org/images/emoji/apple/small_red_triangle.png?v=12" title=":small_red_triangle:" class="emoji" alt=":small_red_triangle:">)</strong> <strong>Conv</strong> <strong>(<img src="https://discuss.pytorch.org/images/emoji/apple/small_red_triangle_down.png?v=12" title=":small_red_triangle_down:" class="emoji" alt=":small_red_triangle_down:">)</strong> <strong>Deconv</strong> <strong>(<img src="https://discuss.pytorch.org/images/emoji/apple/small_red_triangle_down.png?v=12" title=":small_red_triangle_down:" class="emoji" alt=":small_red_triangle_down:">)</strong> <strong>Bilinear</strong> <strong>(<img src="https://discuss.pytorch.org/images/emoji/apple/small_red_triangle_down.png?v=12" title=":small_red_triangle_down:" class="emoji" alt=":small_red_triangle_down:">)</strong> (Benchmark=True, AMP=False, memory_format=channels_last)</p>
<pre><code class="lang-bash">[-------------------------------------- upsample module comparison -------------------------------------]
                          |  nearest(scale=2)  |  conv(scale=3)  |  deconv(scale=2)  |  bilinear(scale=2)
4 threads: ----------------------------------------------------------------------------------------------
      [8, 2048, 16, 16]   |        246.2       |      104.8      |        495.6      |         332.0     
      [16, 2048, 16, 16]  |        486.4       |      197.8      |        973.1      |         665.7     
      [32, 2048, 16, 16]  |        963.8       |      406.1      |       1820.7      |        1328.6     
      [64, 2048, 16, 16]  |       1954.1       |      834.5      |       3646.1      |        2682.6     

Times are in microseconds (us).
</code></pre>
</li>
</ol>
<h3>
<a name="for-those-who-want-to-experiment-with-other-modules-2" class="anchor" href="#for-those-who-want-to-experiment-with-other-modules-2"></a>For those who want to experiment with other modules</h3>
<p>This is script that I used to compare latency.</p>
<pre><code class="lang-auto">import os

os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
os.environ['CUDA_VISIBLE_DEVICES'] = '4'

import torch
import torch.nn as nn
import torch.utils.benchmark as benchmark
import torchvision.models as models
from itertools import product


def get_upsample_module(mode, upsample, ch):
    if mode == 'deconv':
        return nn.Sequential(
            torch.nn.ConvTranspose2d(ch, ch, upsample, stride=upsample, dilation=1, groups=ch, bias=False),
            torch.nn.BatchNorm2d(ch)
        ).cuda()
    else:
        return nn.Upsample(scale_factor=upsample, mode=mode).cuda()

def get_downsample_module(kernel, ch):
    return nn.Sequential(
        torch.nn.Conv2d(ch, ch, kernel, stride=1, padding=1, dilation=1, groups=ch, bias=False),
        torch.nn.BatchNorm2d(ch)
    ).cuda()

use_benchmark = True
amp = False
channel_last = True

torch.backends.cudnn.benchmark = use_benchmark

results = []

batch_size = [8, 16, 32, 64]
channel_size = [2048]
image_size = [16]
mode = ['nearest', 'conv', 'deconv', 'bilinear']
scale_factors = [2]

for b, c, n in product(batch_size, channel_size, image_size):
    label = 'upsample module comparison'
    sub_label = f'[{b}, {c}, {n}, {n}]'
    x = torch.rand((b, c, n, n)).cuda()

    for method, upsample in product(mode, scale_factors):
        if method == 'conv':
            upsample += 1
            model = get_downsample_module(upsample, c)
        else:
            model = get_upsample_module(method, upsample, c)

        if channel_last:
            x = x.to(memory_format=torch.channels_last)
            model = model.to(memory_format=torch.channels_last)
        
        with torch.autocast('cuda', amp):
            results.append(benchmark.Timer(
                stmt='model(x)',
                setup='from __main__ import model',
                globals={'x': x},
                label=label,
                sub_label=sub_label,
                description=f"{method}(scale={upsample})",
                num_threads=4
            ).blocked_autorange(min_run_time=1))

compare = benchmark.Compare(results)
compare.colorize()
compare.print()
</code></pre>