保留正文与补采附件。源页面当前未重新核验；具体限制见交付清单。

Revisit consent button We value your privacy We use cookies to enhance your browsing experience, serve personalised ads or content, and analyse our traffic. By clicking "Accept All", you consent to our use of cookies. Customise Reject All Accept All Customise Consent Preferences We use cookies to help you navigate efficiently and perform certain functions. You will find detailed information about all cookies under each consent category below. The cookies that are categorised as "Necessary" are stored on your browser as they are essential for enabling the basic functionalities of the site. ... Show more NecessaryAlways Active Necessary cookies are required to enable the basic features of this site, such as providing secure log-in or adjusting your consent preferences. These cookies do not store any personally identifiable data.

Functional Functional cookies help perform certain functionalities like sharing the content of the website on social media platforms, collecting feedback, and other third-party features.

Analytics Analytical cookies are used to understand how visitors interact with the website. These cookies help provide information on metrics such as the number of visitors, bounce rate, traffic source, etc.

Performance Performance cookies are used to understand and analyse the key performance indexes of the website which helps in delivering a better user experience for the visitors.

Advertisement Advertisement cookies are used to provide visitors with customised advertisements based on the pages you visited previously and to analyse the effectiveness of the ad campaigns.

Uncategorised Other uncategorised cookies are those that are being analysed and have not been classified into a category as yet. No cookies to display. Reject All Save My Preferences Accept All Toggle SearchToggle Mobile Navigation Publish AI, ML & data-science insights to a global community of data professionals. Sign inSubmit an Article LatestEditor's PicksNewsletterShipAI

Write For TDS LinkedInXToggle Search Deep Learning

Overcoming the Hidden Performance Traps of Variable-Shaped Tensors: Efficient Data Sampling in PyTorch

PyTorch Model Performance Analysis and Optimization — Part 11 Chaim Rand December 3, 20259 min read Photo by National Cancer Institute on Unsplash This is the eleventh part of a series of posts on the topic of analyzing and optimizing PyTorch models. Throughout the series, we have advocated for using the PyTorch Profiler in AI model development and demonstrated the potential impact of performance optimization on the speed and cost of running AI/ML workloads. One common phenomenon we have seen is how seemingly innocent code can hamper runtime performance. In this post, we explore some of the penalties associated with the naive use of variable-shaped tensors — tensors whose shape is dependent on preceding computations and/or inputs. While not applicable to all situations, there are times when the use of variable-shaped tensors can be avoided — although this may come at the expense of additional compute and/or memory. We will demonstrate the tradeoffs of these alternatives on a toy implementation of data sampling in PyTorch.

Three Downsides of Variable Shaped Tensors

We motivate the discussion by presenting three disadvantages to the use of variable-shaped tensors: Learn this step by step with the interactive Machine Learning roadmap.

Host-Device Sync Events

In an ideal scenario, the CPU and GPU are able to run in parallel in an asynchronous manner, with the CPU continuously feeding the GPU with input samples, allocating required GPU memory, and loading GPU compute kernels, and the GPU executing the loaded kernels on the provided inputs using the allocated memory. The presence of dynamic-shaped tensors throws a wrench into this parallelism. In order to allocate the appropriate amount memory, the CPU must wait for the GPU to report the tensor’s shape, and then the GPU must wait for the CPU to allocate the memory and proceed with the kernel loading. The overhead of this sync event can cause a drop in the GPU utilization and slow runtime performance. We saw an example of this in part three of this series when we studied a naive implementation of the common cross-entropy loss that included calls to torch.nonzero and torch.unique. Both APIs return tensors with shapes that are dynamic and dependent on the contents of the input. When these functions are run on the GPU, a host-device synchronization event occurs. In the case of the cross-entropy loss, we discovered the inefficiency through the use of PyTorch Profiler and were able to easily overcome it with an alternative implementation that avoided the use of variable-shaped tensors and demonstrated much better runtime performance.

Graph Compilation

In a recent post we explored the performance benefits of applying just-in-time (JIT) compilation using the torch.compile operator. One of our observations was that graph compilation provided much better results when the graph was static. The presence of dynamic shapes in the graph limits the extent of the optimization via compilation: In some cases, it fails completely; in others it results in lower performance gains. The same implications also apply to other forms of graph compilation, such as XLA, ONNX, OpenVINO, and TensorRT.

Data Batching

Another optimization we have encountered in several of our posts (e.g., here) is sample-batching. Batching improves performance in two primary ways:

  1. Reducing overhead of kernel loading : Rather than loading the GPU kernels required for the computation pipeline once per input sample, the CPU can load the kernels once per batch.
  2. Maximizing parallelization across compute units : GPUs are highly parallel compute engines. The more we are able to parallelize computation, the more we can saturate the GPU and increase its utilization. By batching we can potentially increase the degree of parallelization by a factor of the batch size.

Despite their downsides, the use of variable-shaped tensors is often unavoidable. But sometimes we can modify our model implementation to circumvent them. Sometimes these changes will be straightforward (as in the cross-entropy loss example). Other times they may require some creativity in coming up with a different sequence of fixed-shape PyTorch APIs that provide the same numerical result. Often, this effort can deliver meaningful rewards in runtime and costs. In the next sections, we will study the use of variable-shaped tensors in the context of the data sampling operation. We will start with a trivial implementation and analyze its performance. We will then propose a GPU-friendly alternative that avoids the use of variable-shaped tensors. To compare our implementations, we will use an Amazon EC2 g6e.xlarge with an NVIDIA L40S running an AWS Deep Learning AMI (DLAMI) with PyTorch (2.8). The code we will share is intended for demonstration purposes. Please do not rely on it for accuracy or optimality. Please do not interpret our mention of any framework, library, or platform and an endorsement of its use.

Sampling in AI Model Workloads

In the context of this post, sampling refers to the selection of a subset of items from a large set of candidates for the purposes of computational efficiency, balancing of datatypes, or regularization. Sampling is common in many AI/ML models, such as detection, ranking, and contrastive learning systems. We define a simple variation of the sampling problem: Given a list of N tensors each with a binary label, we are asked to return a subset of K tensors containing both positive and negative examples, in random order. If the input list contains enough samples of each label (K/2), the returned subset should be evenly split. If it is lacking samples of one type, these should be filled with random samples of the second type. The code block below contains a PyTorch implementation of our sampling function. The implementation is inspired by the popular Detectron2 library (e.g., see here and here). For the experiments in this post, we will fix the sampling ratio to 1:10. python

import torch

INPUT_SAMPLES = 10000
SUB_SAMPLE = INPUT_SAMPLES // 10
FEATURE_DIM = 16

def sample_data(input_array, labels):
    device = labels.device
    positive = torch.nonzero(labels == 1, as_tuple=True)[0]
    negative = torch.nonzero(labels == 0, as_tuple=True)[0]
    num_pos = min(positive.numel(), SUB_SAMPLE//2)
    num_neg = min(negative.numel(), SUB_SAMPLE//2)
    if num_neg < SUB_SAMPLE//2:
        num_pos = SUB_SAMPLE - num_neg
    elif num_pos < SUB_SAMPLE//2:
        num_neg = SUB_SAMPLE - num_pos

    # randomly select positive and negative examples
    perm1 = torch.randperm(positive.numel(), device=device)[:num_pos]
    perm2 = torch.randperm(negative.numel(), device=device)[:num_neg]

    pos_idxs = positive[perm1]
    neg_idxs = negative[perm2]

    sampled_idxs = torch.cat([pos_idxs, neg_idxs], dim=0)
    rand_perm = torch.randperm(SUB_SAMPLE, device=labels.device)
    sampled_idxs = sampled_idxs[rand_perm]
    return input_array[sampled_idxs], labels[sampled_idxs]

Performance Analysis With PyTorch Profiler

Even when not immediately obvious, the use of dynamic shapes is easily identifiable in the PyTorch Profiler Trace view. We use the following function to enable PyTorch Profiler: python

def profile(fn, input, labels):
    
    def export_trace(p):
        p.export_chrome_trace(f"{fn.__name__}.json")
        
    with torch.profiler.profile(
            activities=[torch.profiler.ProfilerActivity.CPU,
                        torch.profiler.ProfilerActivity.CUDA],
            with_stack=True,
            schedule=torch.profiler.schedule(wait=0, warmup=10, active=5),
            on_trace_ready=export_trace
    ) as prof:
        for _ in range(20):
            fn(input, labels)
            torch.cuda.synchronize()  # explicit sync for trace readability
            prof.step()

# create random input
input_samples = torch.randn((INPUT_SAMPLES, FEATURE_DIM), device='cuda')
labels = torch.randint(0, 2, (INPUT_SAMPLES,), 
                       device='cuda', dtype=torch.int64)

# run with profiler
profile(sample_data, input_samples, labels)

The image below was captured for the value of ten million input samples. It clearly shows the presence of sync events coming from the torch.nonzero call, as well as the corresponding drops in GPU utilization: Profiler Trace of Sampler (by Author) The use of torch.nonzero in our implementation is not ideal, but can it be avoided?

A GPU-Friendly Data Sampler

We propose an alternative implementation of our sampling function that replaces the dynamic torch.nonzero function with a creative combination of the static torch.count_nonzero, torch.topk, and other APIs: python

def opt_sample_data(input, labels):
    pos_mask = labels == 1
    neg_mask = labels == 0
    num_pos_idxs = torch.count_nonzero(pos_mask, dim=-1)
    num_neg_idxs = torch.count_nonzero(neg_mask, dim=-1)
    half_samples = labels.new_full((), SUB_SAMPLE // 2)
    num_pos = torch.minimum(num_pos_idxs, half_samples)
    num_neg = torch.minimum(num_neg_idxs, half_samples)
    num_pos = torch.where(
        num_neg < SUB_SAMPLE // 2,
        SUB_SAMPLE - num_neg,
        num_pos
    )
    num_neg = SUB_SAMPLE - num_pos

    # create random ordering on pos and neg entries
    rand = torch.rand_like(labels, dtype=torch.float32)
    pos_rand = torch.where(pos_mask, rand, -1)
    neg_rand = torch.where(neg_mask, rand, -1)

    # select top pos entries and invalidate others
    # since CPU doesn't know num_pos, we assume maximum to avoid sync
    top_pos_rand, top_pos_idx = torch.topk(pos_rand, k=SUB_SAMPLE)
    arange = torch.arange(SUB_SAMPLE, device=labels.device)
    if num_pos.numel() > 1:
        # unsqueeze to support batched input
        arange = arange.unsqueeze(0)
        num_pos = num_pos.unsqueeze(-1)
        num_neg = num_neg.unsqueeze(-1)
    top_pos_rand = torch.where(arange >= num_pos, -1, top_pos_rand)

    # repeat for neg entries
    top_neg_rand, top_neg_idx = torch.topk(neg_rand, k=SUB_SAMPLE)
    top_neg_rand = torch.where(arange >= num_neg, -1, top_neg_rand)

    # combine and mix together positive and negative idxs
    cat_rand = torch.cat([top_pos_rand, top_neg_rand], dim=-1)
    cat_idx = torch.cat([top_pos_idx, top_neg_idx], dim=-1)
    topk_rand_idx = torch.topk(cat_rand, k=SUB_SAMPLE)[1]
    sampled_idxs = torch.gather(cat_idx, dim=-1, index=topk_rand_idx)
    sampled_input = torch.gather(input, dim=-2, 
                                 index=sampled_idxs.unsqueeze(-1))
    sampled_labels = torch.gather(labels, dim=-1, index=sampled_idxs)
    return sampled_input, sampled_labels

Clearly, this function requires more memory and more operations than our first implementation. The question is: Do the performance benefits of a static, synchronization-free implementation outweigh the extra cost in memory and compute? To assess the tradeoffs between the two implementations, we introduce the following benchmarking utility: python

def benchmark(fn, input, labels):
    # warm-up
    for _ in range(20):
        _ = fn(input, labels)

    iters = 100
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)
    torch.cuda.synchronize()
    start.record()
    for _ in range(iters):
        _ = fn(input, labels)
    end.record()
    torch.cuda.synchronize()
    avg_time = start.elapsed_time(end) / iters
    
    print(f"{fn.__name__} average step time: {(avg_time):.4f} ms")

benchmark(sample_data, input_samples, labels)
benchmark(opt_sample_data, input_samples, labels)

The following table compares the average runtime of each of the implementations for a variety of input sample sizes: Comparative Step Time Performance — Lower is Better (by Author) For most of the input sample sizes, the overhead of the host-device sync event is either comparable or lower than the additional compute of the static implementation. Disappointingly, we only see a major benefit from the sync-free alternative when the input sample size reaches ten million. Sample sizes that large are uncommon in AI/ML settings. But it’s not our tendency to give up so easily. As noted above, the static implementation enables other optimizations like graph compilation and input batching.

Graph Compilation

Contrary to the original function — which fails to compile — our static implementation is fully compatible with torch.compile: python

benchmark(torch.compile(opt_sample_data), input_samples, labels)

The following table includes the runtimes of our compiled function: Comparative Step Time Performance — Lower is Better (by Author) The results are significantly better — providing a 70–75 percent boost over the original sampler implementation in the 1–10 thousand range. But we still have one more optimization up our sleeve.

Maximizing Performance with Batched Input

Because the original implementation contains variable-shaped operations, it cannot handle batched input directly. To process a batch, we have no choice but to apply it to each input individually, in a Python loop: python

BATCH_SIZE = 32

def batched_sample_data(inputs, labels):
    sampled_inputs = []
    sampled_labels = []
    for i in range(inputs.size(0)):
        inp, lab = sample_data(inputs[i], labels[i])
        sampled_inputs.append(inp)
        sampled_labels.append(lab)
    return torch.stack(sampled_inputs), torch.stack(sampled_labels)

In contrast, our optimized function supports batched inputs as is — no changes necessary. python

input_batch = torch.randn((BATCH_SIZE, INPUT_SAMPLES, FEATURE_DIM),
                          device='cuda')
labels = torch.randint(0, 2, (BATCH_SIZE, INPUT_SAMPLES),
                       device='cuda', dtype=torch.int64)

benchmark(batched_sample_data, input_batch, labels)
benchmark(opt_sample_data, input_batch, labels)
benchmark(torch.compile(opt_sample_data), input_batch, labels)

The table below compares the step times of our sampling functions on a batch size of 32: Step Time Performance on Batched Input — Lower is Better (by Author) Now the results are definitive: By using a static implementation of the data sampler, we are able to boost performance by 2X–52X(!!) the variable-shaped option, depending on the input sample size. Note that although our experiments were run on a GPU device, the model compilation and input batching optimizations also apply to a CPU environment. Thus, avoiding variable shapes could have implications on AI/ML model performance on CPU, as well.

Summary

The optimization process we demonstrated in this post generalizes beyond the specific case of data sampling:

Naturally, not all stories will end as happily as ours. In many cases, we may come across PyTorch code that performs poorly on the GPU but does not have an alternative implementation, or it may have one that requires significantly more compute resources. However, given the potential for meaningful gains in performance and reductions in cost, the process of identifying runtime inefficiencies and exploring alternative implementations is an essential part of AI/ML development.

Register Now to Continue Reading

Unlock Towards Data Science with a Free Account

All articles on Towards Data Science are available to registered users. Create a free account today to access this article and thousands more. Register for freeSign in to your account

Written By

Chaim Rand

See all from Chaim Rand Editor's PicksArtificial IntelligenceData SamplingMachine LearningPytorch Share This Article

Towards Data Science is a community publication. Submit your insights to reach our global audience and earn through the TDS Author Payment Program. Write for TDS

Related Articles

From Resume to Cover Letter Using AI and LLM, with Python and Streamlit

Large Language Models Here’s how less than 100 lines of code can build you an LLM app that converts your resume into your job-specific cover letter Piero Paialunga February 5, 202513 min read

Your Vision-Language Model Might Be a Bag of Words

Artificial Intelligence We explore the limits of what vision-language models get about language in our Oral Paper at ICLR 2023 Federico Bianchi March 21, 20239 min read

Hyperparameter Optimization - Intro and Implementation of Grid Search, Random Search and Bayesian...

Artificial Intelligence Most common hyperparameter optimization methodologies to boost machine learning outcomes. Farzad Nobar March 13, 202310 min read

Analysis of US Elections from 1976 to 2010 with Pandas

Artificial Intelligence Practice Pandas while providing valuable insights into past elections. Soner Yıldırım October 18, 20205 min read

8 Must-Have Git Commands for Data Scientists

Artificial Intelligence This is how you contribute to a project Soner Yıldırım July 18, 20214 min read

Why Your AI Search Evaluation Is Probably Wrong (And How to Fix It)

Artificial Intelligence A five-step framework for building rigorous, reproducible AI search benchmarks — before you make six-figure infrastructure decisions Zairah Mustahsan March 9, 20266 min read

Fit vs. Transform in SciKit libraries for Machine Learning

Machine Learning We have seen methods such as fit(), transform(), and fit_transform() in a lot of SciKit's libraries. And almost all tutorials, including... Sunny Srinidhi December 3, 20193 min read

5 Tips That Will Help You Become a Data Science Manager

Data Science So you want to be promoted? So does everyone else. How can you stand out from your peers? Tessa Xie May 24, 20226 min read

Intergenerational Mobility in the US - A Data Science Problem (1/5)

Data Science This series was written together with Ibukun Aribilola and Valdrin Jonuzi as part of the Data Science for Social Good tutorial. Ang Li-Lian April 22, 202211 min read

Your home for data science and AI. The world's leading publication for data science, data analytics, data engineering, machine learning, and artificial intelligence professionals. Subscribe to Our NewsletterWrite For TDS•About•Advertise•Privacy Policy•Terms of Use © Insight Media Group, LLC 2026•The New Stack•roadmap.sh•Creator Network

Recommended Articles

Close