Why can‘t pytorch profiler capture activities of data preprocess in __getitem__()?

username: xiaoyeguan6
cooked: <h2><a name="profiling-data-preprocessing-activities-in-__getitem__-1" class="anchor" href="#profiling-data-preprocessing-activities-in-__getitem__-1"></a>Profiling data preprocessing activities in <strong>getitem</strong></h2>
<p>Hi all,</p>
<p>I’m currently using <code>torch.profiler</code> to profile the performance of my data loading process. However, I’m running into an issue where the profiler seems to aggregate the entire data loading process under a single <code>enumerate(dataloader)</code> operation, rather than breaking it down into the more granular <code>record_function</code> calls that I’ve placed inside my <code>__getitem__</code> method.</p>
<p>Here’s the relevant part of my code:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
from torch.utils.data import Dataset, DataLoader
from torch.profiler import profile, record_function, ProfilerActivity

class KiTSDataset(Dataset):
    def __init__(self, cases, crop_size=None, fixed_depth=None):
        self.cases = cases
        self.crop_size = crop_size
        self.fixed_depth = fixed_depth
    
    def __len__(self):
        return len(self.cases)

    def crop_center(self, volume, crop_size):
        _, d, h, w = volume.shape
        start_d = (d - crop_size[0]) // 2
        start_h = (h - crop_size[1]) // 2
        start_w = (w - crop_size[2]) // 2
        return volume[:, start_d:start_d + crop_size[0], start_h:start_h + crop_size[1], start_w:start_w + crop_size[2]]

    def pad_or_crop_depth(self, volume):
        depth = volume.shape[1]
        if depth &lt; self.fixed_depth:
            pad_size = self.fixed_depth - depth
            volume = F.pad(volume, (0, 0, 0, 0, pad_size // 2, pad_size - pad_size // 2))
        elif depth &gt; self.fixed_depth:
            start = (depth - self.fixed_depth) // 2
            volume = volume[:, start:start + self.fixed_depth, :, :]
        return volume

    def __getitem__(self, idx):
        case = self.cases[idx]
        with record_function("loadcase"):
            volume, segmentation = load_case(case)
        
        with record_function("ToTensor"):
            volume = torch.tensor(volume.get_fdata(), dtype=torch.float32).unsqueeze(0)
            segmentation = torch.tensor(segmentation.get_fdata(), dtype=torch.long).unsqueeze(0)

        if self.crop_size:
            with record_function("CropCenter"):
                volume = self.crop_center(volume, self.crop_size)
                segmentation = self.crop_center(segmentation, self.crop_size)

        if self.fixed_depth:
            with record_function("PadOrCropDepth"):
                volume = self.pad_or_crop_depth(volume)
                segmentation = self.pad_or_crop_depth(segmentation)

        return volume, segmentation.squeeze(0)

def main():
    cases = [f'case_{i:05d}' for i in range(100)]
    dataset = KiTSDataset(cases, crop_size=(100, 256, 256), fixed_depth=10)
    dataloader = DataLoader(dataset, batch_size=2, shuffle=True, num_workers=2)

    with profile(
        activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
        record_shapes=True,
        profile_memory=True,
        with_stack=True
    ) as prof:
        for data in dataloader:
            with record_function("model_inference"):
                # Simulate model inference
                pass
        prof.step()

    print(prof.key_averages().table(sort_by="self_cpu_time_total"))

if __name__ == '__main__':
    main()
</code></pre>
<p>Here is a snapshot of the profiler’s output:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/4/d4847fd1a5b77e212748d05a14512921680c9fd6.png" data-download-href="https://discuss.pytorch.org/uploads/default/d4847fd1a5b77e212748d05a14512921680c9fd6" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/d/4/d4847fd1a5b77e212748d05a14512921680c9fd6_2_690x206.png" alt="image" data-base62-sha1="uk16kHJrOrlLvtNMmlPQO08anjM" width="690" height="206" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/d/4/d4847fd1a5b77e212748d05a14512921680c9fd6_2_690x206.png, https://discuss.pytorch.org/uploads/default/optimized/3X/d/4/d4847fd1a5b77e212748d05a14512921680c9fd6_2_1035x309.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/d/4/d4847fd1a5b77e212748d05a14512921680c9fd6_2_1380x412.png 2x" data-dominant-color="232323"><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">1846×553 56.7 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>

username: xiaoyeguan6
cooked: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> Hi, can you please help me out here. Or is there possible to use nsys to profile every activity of data preprocess?</p>

username: ptrblck
cooked: <p>I don’t think the profiler works in a multi-processing environment and you might need to use <code>num_workers=0</code> to profile the data loading pipeline.<br>
Here is a simple example showing this behavior:</p>
<pre data-code-wrap="python"><code class="lang-python">class MyDataset(Dataset):
    def __init__(self):
        self.data = torch.randn(10, 3, 224, 224)
        self.target = torch.randint(0, 10, (10,))
        
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, index):
        x = self.data[index]
        y = self.target[index]
        
        with record_function("transform1"):
                x = x * 2
                
        with record_function("transform2"):
                y = y + 1
        
        return x, y
    
dataset = MyDataset()
loader = DataLoader(dataset, batch_size=5, num_workers=0)

with profile(activities=[ProfilerActivity.CPU], record_shapes=True) as prof:
    with record_function("loader"):
        for data, target in loader:
            pass

print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=15))
# -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  
#                                                    Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg    # of Calls  
# -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  
#                                                  loader        10.30%     306.577us       100.00%       2.978ms       2.978ms             1  
# enumerate(DataLoader)#_SingleProcessDataLoaderIter._...         8.93%     266.004us        89.21%       2.657ms     885.566us             3  
#                                              transform1         6.86%     204.200us        56.38%       1.679ms     167.883us            10  
#                                               aten::mul        47.42%       1.412ms        49.52%       1.475ms     147.463us            10  
#                                             aten::stack         0.70%      20.700us        16.14%     480.542us     120.135us             4  
#                                               aten::cat        14.62%     435.316us        15.04%     447.960us     111.990us             4  
#                                              transform2         4.75%     141.413us         5.67%     168.766us      16.877us            10  
#                                                aten::to         0.30%       8.915us         2.10%      62.627us       6.263us            10  
#                                            aten::select         1.73%      51.482us         2.10%      62.556us       3.128us            20  
#                                          aten::_to_copy         0.81%      24.018us         1.80%      53.712us       5.371us            10  
#                                               aten::add         0.92%      27.353us         0.92%      27.353us       2.735us            10  
#                                             aten::copy_         0.68%      20.229us         0.68%      20.229us       2.023us            10  
#                                        aten::as_strided         0.48%      14.319us         0.48%      14.319us       0.447us            32  
#                                            aten::narrow         0.20%       5.941us         0.42%      12.644us       6.322us             2  
#                                     aten::empty_strided         0.32%       9.465us         0.32%       9.465us       0.947us            10  
# -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  
# Self CPU time total: 2.978ms
</code></pre>

username: xiaoyeguan6
cooked: <p>Thanks, that helps a lot.</p>