Keypoint RCNN training time increases per epoch

blackmamba: <p>Hi,</p>
<p>I am training a keypoint rcnn model and noticed that on every epoch, the training time increases.</p>
<p>For the first epoch, it takes about 32 seconds, then the next takes about 45, then it goes upto 60 and increases about a second every two epochs.</p>
<p>Just wanted to know if this a normal behaviour for this model.</p>
<p>Using the latest pytorch nightly on cuda 12.1. I did test on the stable release as well with cuda 11.7 and 11.8 but noticed the same issues.</p>
<p>Running on a 40 series laptop GPU with nvidia 535 drivers</p>

ptrblck: <p>To narrow down the issue properly you should profile the workload and check what exactly slows down during the later epochs. In case you are also seeing and increase in device memory usage in each iteration, check if you are storing model outputs or the loss in a <code>list</code> without detaching it as it would also store the entire computation graph.<br>
If that’s not the case, my guess is that your laptop might be overheating and could throttle some clocks which might slow down the training.</p>

blackmamba: <p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/3/1/31aca7bfdc9299bc0838284eaa6198488966f7d3.png" data-download-href="https://discuss.pytorch.org/uploads/default/31aca7bfdc9299bc0838284eaa6198488966f7d3" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/3/1/31aca7bfdc9299bc0838284eaa6198488966f7d3_2_690x362.png" alt="image" data-base62-sha1="75rj6Y8isTTITEvC78FVmunaQ6L" width="690" height="362" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/3/1/31aca7bfdc9299bc0838284eaa6198488966f7d3_2_690x362.png, https://discuss.pytorch.org/uploads/default/optimized/3X/3/1/31aca7bfdc9299bc0838284eaa6198488966f7d3_2_1035x543.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/3/1/31aca7bfdc9299bc0838284eaa6198488966f7d3_2_1380x724.png 2x" data-dominant-color="3A162E"><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">1592×836 156 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>This is the result of the profiler. I believe the <code>torch.where</code> is taking most of the cpu time which is present throughtout the Keypoint RCNN torch model.</p>
<p>But why would it get slower per epoch/optimization step?</p>

ptrblck: <p>Compare profiles of a fast vs. slow epoch and check where the difference is. I would also recommend using a visual profiler, such as Nsight Systems, to see the real timeline instead of a summary of the kernels.</p>

blackmamba: <p>Hi,</p>
<p>So I tried using NSight and found that there are long paues for <code>cudaStreamSynchronize</code> while loading the training images. The image size is ‘1536,1536’, I push them to GPU and it gets resized internally to 512,512 in Keypoint RCNN.</p>
<p>I tried a few things:</p>
<ol>
<li>Resizing the image to 256,256 before passing it to keypoint RCNN. Note that the all the target data is wrong since everything is calcluated with respect to the 1536x1536 image.</li>
</ol>
<p>In this scenario, I get a constant 14-15 seconds training time per epoch.</p>
<ol start="2">
<li>
<p>Using <code>non-blocking = True</code> with 256x256, this remains the same as in 1.</p>
</li>
<li>
<p>Using <code>non_blocking=True</code> with 1536x1536 image sizes, this looks like it reduces the training time a bit but there is still an increase in training time per epoch, it goes from 26 seconds on the first epoch, to 32 seconds, then on the 6th epoch its on 46 seconds.</p>
</li>
</ol>
<p>Using non-blocking the data loads really quickly, however the <code>cudaStreamSynchronize</code> call moves into the <code>model()</code> call of keypoint rcnn</p>

ptrblck: <p>It seems your data loading might be the bottleneck. Did you check where the sync is coming from in the <code>nsys</code> profile?<br>
Using <code>non_blocking=True</code> would allow the CPU to run ahead and not to block on the data transfer. However, if the next operation (e.g. the <code>forward</code> method) requires the data on the GPU you would then still need to wait until the transfer is done before the kernel can be launched.</p>

blackmamba: <p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/7/c/7ce320ccd33c806c3678430e6a38baa89f24936c.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/7ce320ccd33c806c3678430e6a38baa89f24936c" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/7/c/7ce320ccd33c806c3678430e6a38baa89f24936c_2_690x347.jpeg" alt="image" data-base62-sha1="hONPdwLUZ9gdw7MRpfel4mSqHms" width="690" height="347" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/7/c/7ce320ccd33c806c3678430e6a38baa89f24936c_2_690x347.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/7/c/7ce320ccd33c806c3678430e6a38baa89f24936c_2_1035x520.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/7/c/7ce320ccd33c806c3678430e6a38baa89f24936c_2_1380x694.jpeg 2x" data-dominant-color="4D514D"><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×968 126 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>Sorry, I am not sure if I followed it right. Is this the way to check it?</p>

ptrblck: <p>Yes, as the stacktrace shows the sync is triggered by the copy kernel, so indeed your model seems to wait for the data transfer.</p>

blackmamba: <p>Does that mean I am bound by my hardware?</p>
<p>Also, why would it get slower over time?</p>

blackmamba: <p>Hi</p>
<p>I did some more testing and found that the code slows down only when I call the <code>backward</code> function.</p>
<p>Loading the data and calling <code>forward</code> seem to be fast enough. When I introduce the backward call, the loader seems to get slower over time</p>

ptrblck: <p>Did you compare the profiles of both runs and checked if the sync caused by the data loading is still the bottleneck? I wouldn’t know how the <code>backward</code> call could slow down your data loading, so a profile might again help here.</p>

blackmamba: <pre><code class="lang-auto">    def train_one_epoch(self, train=True):
        self.model.train()
        loader = self.train_loader if train else self.val_loader
        with torch.set_grad_enabled(train):
            for images, targets in loader:
                torch.cuda.nvtx.range_push("Gradient Step")
                torch.cuda.nvtx.range_push("Loading Data")
                torch.cuda.nvtx.range_push("Loading Images")
                images = [image.to(self.device) for image in images]
                torch.cuda.nvtx.range_pop()
                torch.cuda.nvtx.range_push("Loading Targets")
                targets = [
                    {k: v.to(self.device) for k, v in t.items()} for t in targets
                ]
                torch.cuda.nvtx.range_pop()
                torch.cuda.nvtx.range_pop()
                torch.cuda.nvtx.range_push("Forward Pass")
                loss_dict = self.model(images, targets)
                torch.cuda.nvtx.range_pop()
                losses = sum(loss for loss in loss_dict.values())
                if train:
                    torch.cuda.nvtx.range_push("Backward Pass")
                    self.optimizer.zero_grad()
                    losses.backward()
                    self.optimizer.step()
                    torch.cuda.nvtx.range_pop()
                torch.cuda.nvtx.range_pop()
</code></pre>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/1/f/1fb4ded0bbc53b9c281ab0a6ea65e1b3b7ce2356.png" data-download-href="https://discuss.pytorch.org/uploads/default/1fb4ded0bbc53b9c281ab0a6ea65e1b3b7ce2356" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/1/f/1fb4ded0bbc53b9c281ab0a6ea65e1b3b7ce2356_2_690x186.png" alt="image" data-base62-sha1="4wujvcNRsq1eiDW1QbkBKqpfqcK" width="690" height="186" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/1/f/1fb4ded0bbc53b9c281ab0a6ea65e1b3b7ce2356_2_690x186.png, https://discuss.pytorch.org/uploads/default/optimized/3X/1/f/1fb4ded0bbc53b9c281ab0a6ea65e1b3b7ce2356_2_1035x279.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/1/f/1fb4ded0bbc53b9c281ab0a6ea65e1b3b7ce2356_2_1380x372.png 2x" data-dominant-color="444642"><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">3207×867 137 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>The above screenshot is obatined by omitting <code>losses.backward()</code> and <code>optimizer.step()</code>. There is very little time spent on synchronize during the data loading process.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/8/b8d4b39ac3b918259a7e518e42215eab02484a02.png" data-download-href="https://discuss.pytorch.org/uploads/default/b8d4b39ac3b918259a7e518e42215eab02484a02" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/b/8/b8d4b39ac3b918259a7e518e42215eab02484a02_2_690x186.png" alt="image" data-base62-sha1="qn5z49vWiTtFT4ZZJNiq0xgOudk" width="690" height="186" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/b/8/b8d4b39ac3b918259a7e518e42215eab02484a02_2_690x186.png, https://discuss.pytorch.org/uploads/default/optimized/3X/b/8/b8d4b39ac3b918259a7e518e42215eab02484a02_2_1035x279.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/b/8/b8d4b39ac3b918259a7e518e42215eab02484a02_2_1380x372.png 2x" data-dominant-color="424642"><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">3207×867 133 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><br>
This is when I have the backward and step running.</p>

ptrblck: <p>Sorry, but the screenshots are not giving me enough information. Could you post a minimal and executable code snippet which would show the behavior so that I could profile it?</p>

blackmamba: <pre><code class="lang-auto">import torch
import torchvision
from torch.utils.data import DataLoader
from torch.utils.data.dataset import Dataset
from torchvision.models import ResNet50_Weights
from torchvision.models.detection.rpn import AnchorGenerator


class FakeDataset(Dataset):
    def __init__(self, num_keypoints, num_classes):
        num_images = 800
        num_instances_per_image = 3
        width, height = 1536, 1536
        self.images = torch.rand(num_images, 1, width, height)
        self.labels = torch.randint(
            0,
            num_classes,
            (
                num_images,
                num_instances_per_image,
            ),
        )
        bboxes_min = torch.rand(num_images, num_instances_per_image, 2) * width // 2
        bboxes_max = bboxes_min + torch.randint(
            1, width // 2, size=(num_images, num_instances_per_image, 2)
        )
        self.bboxes = torch.cat([bboxes_min, bboxes_max], dim=2)
        self.keypoints = (
            torch.rand(num_images, num_instances_per_image, num_keypoints, 2)
            * width
            // 2
        )
        self.keypoints = torch.cat(
            [self.keypoints, torch.ones_like(self.keypoints)], dim=3
        )

    def __getitem__(self, index):
        data_dict = {}
        data_dict["labels"] = self.labels[index]
        data_dict["boxes"] = self.bboxes[index]
        data_dict["keypoints"] = self.keypoints[index]
        return self.images[index], data_dict

    def __len__(self):
        return len(self.images)


def timer(func):
    def wrapper(*args, **kwargs):
        import time

        start = time.time()
        result = func(*args, **kwargs)
        end = time.time()
        print(f"Time taken to run {func.__name__}: {end - start} seconds")
        return result

    return wrapper


class Trainer:
    def __init__(self, model, train_loader, val_loader, config):
        self.model = model.to(config["device"])
        self.train_loader = train_loader
        self.val_loader = val_loader
        self.optimizer = config["optimizer"]
        self.device = config["device"]
        self.num_epochs = config["num_epochs"]
        self.eval_iter = 10

    @timer
    def train_one_epoch(self, train=True):
        self.model.train()

        loader = self.train_loader if train else self.val_loader
        total_loss = 0
        with torch.set_grad_enabled(train):
            for images, targets in loader:
                images = [image.to(self.device) for image in images]
                targets = [
                    {k: v.to(self.device) for k, v in t.items()} for t in targets
                ]
                loss_dict = self.model(images, targets)
                losses = sum(loss for loss in loss_dict.values())
                total_loss += losses.item()
                if train:
                    self.optimizer.zero_grad()
                    losses.backward()
                    self.optimizer.step()
        return total_loss

    def train(self):
        for epoch in range(1, self.num_epochs + 1):
            train_loss = self.train_one_epoch(train=True)
            print("Epoch: {}, Train Loss: {}".format(epoch, train_loss))
            if epoch &gt; 1 and epoch % self.eval_iter == 0:
                val_loss = self.train_one_epoch(train=False)


if __name__ == "__main__":
    anchor_generator = AnchorGenerator(
        sizes=(32, 64, 128, 256, 512),
        aspect_ratios=(0.25, 0.5, 0.75, 1.0, 2.0, 3.0, 4.0),
    )
    model = torchvision.models.detection.keypointrcnn_resnet50_fpn(
        weights=None,
        weights_backbone=ResNet50_Weights.DEFAULT,
        num_keypoints=10,
        # Background is the first class, object is the second class
        num_classes=10,
        rpn_anchor_generator=anchor_generator,
        fixed_size=(512, 512),
    )

    dataset = FakeDataset(num_keypoints=10, num_classes=10)
    data_loader = DataLoader(
        dataset,
        batch_size=8,
        shuffle=True,
        num_workers=2,
        collate_fn=lambda x: tuple(zip(*x)),
    )

    params = [p for p in model.parameters() if p.requires_grad]
    optimizer = torch.optim.SGD(params, lr=0.001, momentum=0.9, weight_decay=0.0005)

    config = {
        "optimizer": optimizer,
        "device": torch.device("cuda" if torch.cuda.is_available() else "cpu"),
        "num_epochs": 30,
    }

    train_loader = data_loader
    val_loader = data_loader

    trainer = Trainer(model, train_loader, val_loader, config)
    trainer.train()


</code></pre>
<p>Added a random dataset too. After every epoch, the training time increases 4 seconds. Should be able to run the script as is</p>

blackmamba: <p>Hi,</p>
<p>Is there any update on this?</p>

blackmamba: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> Sorry, but were you able to reproduce it?</p>