Why backward and optimizer.step become slower with jit.trace

Co1lin: <p>Hi all! I am trying to use JIT to speed up my training. However, when profiling my code with JIT trace, I found that while zero_grad and forward are much more faster than the case without JIT, backward and optimizer.step become slower. Why this happened? I am also wondering whether JIT compiler has some optimizations for backward and optimizer.step. I tried to search on the Internet but found nothing about this.</p>
<p>Profiling result:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/4/9/49f1752094b2de83e820bf0e2e03fbc742ec47a8.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/49f1752094b2de83e820bf0e2e03fbc742ec47a8" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/4/9/49f1752094b2de83e820bf0e2e03fbc742ec47a8_2_690x305.jpeg" alt="image" data-base62-sha1="ay8aTqYjayimyCTxbeqa2PSJYrm" width="690" height="305" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/4/9/49f1752094b2de83e820bf0e2e03fbc742ec47a8_2_690x305.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/4/9/49f1752094b2de83e820bf0e2e03fbc742ec47a8_2_1035x457.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/4/9/49f1752094b2de83e820bf0e2e03fbc742ec47a8_2_1380x610.jpeg 2x" data-dominant-color="343335"><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">1920×850 132 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>Code to reproduce this:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
import torch.nn as nn

from IPython import embed

def profile(func):
    from line_profiler import LineProfiler

    def wrapper(*args, **kwargs):
        lp = LineProfiler()
        lp_wrapper = lp(func)
        result = lp_wrapper(*args, **kwargs)
        lp.print_stats()

        return result
    return wrapper


class MyModule(nn.Module):
    def __init__(self):
        super().__init__()

        self.layers = nn.Sequential(
            nn.Conv2d(3, 64, 3, 1, 1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.Conv2d(64, 64, 3, 1, 1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.Conv2d(64, 32, 3, 1, 1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.AvgPool2d(2, 2),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.Flatten(),
            nn.Linear(32*64*64, 128),
            nn.BatchNorm1d(128),
            nn.ReLU(),
            nn.Linear(128, 8),
        )
        self.loss_fn = nn.CrossEntropyLoss()

    def forward(self, x, target):
        x = self.layers(x)
        return x, self.loss_fn(x, target)


@profile
def perf_train(model, inp, target):
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
    criterion = nn.CrossEntropyLoss()        

    import time
    for _ in range(50): # warmup
        optimizer.zero_grad()
        out, loss = model(inp, target)
        loss.backward()
        optimizer.step()
    start = time.time_ns()
    for _ in range(500):
        optimizer.zero_grad()
        out, loss = model(inp, target)
        loss.backward()
        optimizer.step()
    print(f'dur: {(time.time_ns() - start) / 1e6} ms')
    return out


def eager(perf_func):
    device = torch.device('cuda')

    mod = MyModule().to(device)
    inp = torch.randn(128, 3, 128, 128).to(device)
    target = torch.randint(0, 8, (128,)).to(inp.device)

    out = perf_func(mod, inp, target)
    
    # print(out.shape)


def trace(perf_func):
    device = torch.device('cuda')

    mod = MyModule().to(device)
    inp = torch.randn(128, 3, 128, 128).to(device)
    target = torch.randint(0, 8, (128,)).to(inp.device)

    traced = torch.jit.trace(mod, (inp, target))

    out = perf_func(traced, inp, target)
    
    # print(out.shape)

# run one of them each time
# eager(perf_train)
trace(perf_train)
</code></pre>

ptrblck: <p>It seems you are using the GPU for your model training and <code>LineProfiler</code> to profile the code. I doubt this profiler is aware of CUDA’s asynchronous execution and would claim the results might be invalid or did you correlate them to manual profiling via synchronized timers or via event based profiling?</p>