2x faster training than default precision when model is just 3 linear layer, but when the model is just 1 conv layer, mixed precision training is slower than default precision

hotpanda: <p>OS: ubuntu 16.04<br>
Cuda: 10.2<br>
pytorch: 1.8.0<br>
I ran the demo <a href="https://pytorch.org/tutorials/recipes/recipes/amp_recipe.html#a-simple-network" rel="noopener nofollow ugc">https://pytorch.org/tutorials/recipes/recipes/amp_recipe.html#a-simple-network</a>.<br>
default precision training time is 1.760 sec and mixed precision training time is 0.833. 2x speed up. But when I use torch.nn.Conv2d, amp training is slower than default model. please tell me why?</p>

ptrblck: <p>Could you share the code used to profile these workloads as well as more information about your setup, please? In particular information about the used GPU would be needed.</p>

hotpanda: <pre><code class="lang-auto">import torch
import torch.nn as nn
from torch.cuda.amp import autocast, GradScaler
from torchvision.models import resnet18
from torch.backends.cudnn import benchmark
import time
import gc

start_time = None
torch.backends.cudnn.benchmark = True

batch_size = 512
in_size = 4096
out_size = 4096
num_layers = 3
num_batches = 50
epochs = 3
data = [torch.randn(batch_size, in_size, device='cuda') for _ in range(num_batches)]
targets = [torch.randn(batch_size, out_size, device='cuda') for _ in range(num_batches)]

def start_timer():
    global start_time
    gc.collect()
    torch.cuda.empty_cache()
    torch.cuda.reset_max_memory_allocated()
    torch.cuda.synchronize()
    start_time = time.time()


def end_timer_and_print(local_msg):
    torch.cuda.synchronize()
    end_time = time.time()
    print('\n' + local_msg)
    print('Total execution time = {:.3f} sec'.format(end_time -start_time))
    print('Max memory used by tensors = {} bytes'.format(torch.cuda.max_memory_allocated()))


def make_model(in_size, out_size, num_layers):
    layers = []
    for _ in range(num_layers - 1):
        layers.append(nn.Linear(in_size, in_size))
        layers.append(nn.ReLU())
    layers.append(nn.Linear(in_size, out_size))
    return nn.Sequential(*tuple(layers)).cuda()

def train_default():
    net = make_model(in_size, out_size, num_layers)
    opt = torch.optim.SGD(net.parameters(), lr=0.001)
    loss_fn = torch.nn.MSELoss().cuda()
    net.train()
    start_timer()
    for epoch in range(epochs):
        for input, target in zip(data, targets):
            output = net(input)
            loss = loss_fn(output, target)
            loss.backward()
            opt.step()
            opt.zero_grad()
    end_timer_and_print('Default precision:')

def train_mixed():
    net = make_model(in_size, out_size, num_layers)
    opt = torch.optim.SGD(net.parameters(), lr=0.001)
    loss_fn = torch.nn.MSELoss().cuda()
    net.train()
    scaler = GradScaler()
    start_timer()
    for epoch in range(epochs):
        for input, target in zip(data, targets):
            with autocast():
                output = net(input)
                loss = loss_fn(output, input)
            scaler.scale(loss).backward()
            scaler.step(opt)
            scaler.update()
            opt.zero_grad()
    end_timer_and_print('Mixed Precision:')

class SingleConv(nn.Module):
    def __init__(self):
        super(SingleConv, self).__init__()
        self.conv = nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1, bias=False)

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

def train_single_conv():
    default_net = SingleConv()
    default_net.cuda()
    default_net.train()
    batches = 5
    default_opt = torch.optim.SGD(default_net.parameters(), lr=0.0001)
    loss_fn = torch.nn.MSELoss().cuda()
    data = [torch.randn(64, 3, 224, 224, device='cuda') for _ in range(batches)]
    tar = [torch.randn(64, 64, 112, 112, device='cuda') for _ in range(batches)]
    start_timer()
    for _ in range(5):
        for input, target in zip(data, tar):
            out = default_net(input)
            loss = loss_fn(out, target)
            loss.backward()
            default_opt.step()
            default_opt.zero_grad()
    end_timer_and_print('Default 1 Conv train:')

    mixed_net = SingleConv()
    mixed_net.cuda()
    mixed_net.train()
    loss_fn = nn.MSELoss().cuda()
    scaler = GradScaler()
    mixed_opt = torch.optim.SGD(mixed_net.parameters(), lr=0.0001)
    start_timer()
    for _ in range(5):
        for input, target in zip(data, tar):
            with autocast():
                out = mixed_net(input)
                loss = loss_fn(out, target)
            scaler.scale(loss).backward()
            scaler.step(mixed_opt)
            scaler.update()
            mixed_opt.zero_grad()
    end_timer_and_print('Mixed 1 Conv train:')
</code></pre>
<p>the output is:<br>
Default precision:<br>
Total execution time = 1.760 sec<br>
Max memory used by tensors = 1367458816 bytes</p>
<p>Mixed Precision:<br>
Total execution time = 0.833 sec<br>
Max memory used by tensors = 1401006592 bytes</p>
<p>Default 1 Conv train:<br>
Total execution time = 0.090 sec<br>
Max memory used by tensors = 2470273536 bytes</p>
<p>Mixed 1 Conv train:<br>
Total execution time = 0.127 sec<br>
Max memory used by tensors = 2644665344 bytes</p>

hotpanda: <p>And I used resnet18 to test. amp is faster, the memory usage is also lower.</p>
<pre><code class="lang-auto">def train_resnet18():
    x = torch.randn(64, 3, 224, 224, device='cuda')
    y = torch.randn(64, 1000, device='cuda')
    default_model = resnet18(pretrained=False)
    default_model.cuda()
    default_model.train()
    opt = torch.optim.SGD(default_model.parameters(), lr=1e-3)
    loss_fn = torch.nn.MSELoss().cuda()
    start_timer()
    for _ in range(500):
        y_pred = default_model(x)
        loss = loss_fn(y, y_pred)

        loss.backward()
        opt.step()
        opt.zero_grad()
    end_timer_and_print('default resnet18 training')

    mixed_model = resnet18(pretrained=False)
    mixed_model.cuda()
    mixed_model.train()
    scaler = GradScaler()
    mix_opt = torch.optim.SGD(mixed_model.parameters(), lr=1e-3)
    loss_fn = torch.nn.MSELoss().cuda()
    start_timer()
    for _ in range(500):
        with autocast():
            pred = mixed_model(x)
            loss = loss_fn(y, pred)
        scaler.scale(loss).backward()
        scaler.step(mix_opt)
        scaler.update()
        mix_opt.zero_grad()
    end_timer_and_print('mixed resnet18 training')
</code></pre>
<p>default resnet18 training<br>
Total execution time = 27.581 sec<br>
Max memory used by tensors = 3741790208 bytes</p>
<p>mixed resnet18 training<br>
Total execution time = 15.231 sec<br>
Max memory used by tensors = 3247211520 bytes</p>

ptrblck: <p>Thanks, I still don’t know the device you were using.</p>

hotpanda: <p>sorry, device is v100, 32g</p>

ptrblck: <p>Thanks for the update!<br>
I’ve added some missing warmup iterations, since you are using cuDNN’s benchmark mode, so the first iteration will see a slowdown due to the internal profiling.<br>
I can reproduce the slower kernel execution time in FP16 using the default memory format, but see a speedup using <code>memory_format=torch.channels_last</code>, which will avoid the internal permutations via:</p>
<pre><code class="lang-python">cudnn::ops::nchwToNhwcKernel
cudnn::ops::nhwcToNchwKernel
</code></pre>
<p>Used code:</p>
<pre><code class="lang-python">import torch
import torch.nn as nn
import time

torch.backends.cudnn.benchmark = True

conv = nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1, bias=False).cuda()
x = torch.randn(64, 3, 224, 224).cuda()


# warmup
for _ in range(10):
    out = conv(x)
    out.backward(torch.rand_like(out))
grad = torch.rand_like(out)

nb_iters = 100
torch.cuda.synchronize()
t0 = time.perf_counter()
for _ in range(nb_iters):
    out = conv(x)
    out.backward(grad)
torch.cuda.synchronize()
t1 = time.perf_counter()
print('{} s/iter'.format((t1 - t0)/nb_iters))


# warmup
conv.to(memory_format=torch.channels_last)
x = x.to(memory_format=torch.channels_last)
for _ in range(10):
    with torch.cuda.amp.autocast():
        out = conv(x)
        out.backward(torch.rand_like(out))
grad = torch.rand_like(out)

torch.cuda.synchronize()
t0 = time.perf_counter()
for _ in range(nb_iters):
    with torch.cuda.amp.autocast():
        out = conv(x)
        out.backward(grad)
torch.cuda.synchronize()
t1 = time.perf_counter()
print('{} s/iter'.format((t1 - t0)/nb_iters))
</code></pre>
<p>Default memory format:</p>
<pre><code class="lang-python">0.0012196038290858268 s/iter
0.0014858345501124858 s/iter
</code></pre>
<p>Channels-last:</p>
<pre><code class="lang-python">0.0012191115505993367 s/iter
0.0008443104848265648 s/iter
</code></pre>

hotpanda: <p>Thank you for the reply. I’ll try this later.<br>
And I can not understand clearly about your reply “Grouped convolutions might not trigger the FP16 path and thus might not use TensorCores.<a href="https://discuss.pytorch.org/t/f-conv2d-runs-x15-slower-in-mixed-precision-half-precision-mode/69987">here</a>”. what does the “Grouped convolutions” mean. Could you please give me some example?</p>

ptrblck: <aside class="quote no-group" data-username="hotpanda" data-post="8" data-topic="131291">
<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/hotpanda/48/31596_2.png" class="avatar"> hotpanda:</div>
<blockquote>
<p>what does the “Grouped convolutions” mean. Could you please give me some example?</p>
</blockquote>
</aside>
<p>By grouped convolution I’m meaning a conv layer using <code>groups!=1</code> such as in the linked post:</p>
<pre data-code-wrap="python"><code class="lang-python">out = F.conv2d(input, self.weight, padding=self.padding, groups=4)
</code></pre>

hotpanda: <p>Thanks you so much. this helps me a lot.</p>