How do you optimize inference batch size in c++ pytorch libtorch?

username: kevj
cooked: <p>I have an inference service using cpu (will migrate to gpu later but right now I am stuck with cpu) that is taking request through thrift and returning inference result</p>
<p>Right now its doing inference 1 request at a time, I think libtorch internally have some cpu optimization like SIMD that can make batch inference more optimal</p>
<p>here a simple c++ code I test with a toy model (model that emulate euclidean distance) with benchmark of each batch number</p>
<pre><code>    #include &lt;torch/script.h&gt; // One-stop header.
    #include &lt;iostream&gt;
    #include &lt;chrono&gt;
    #include &lt;vector&gt;
    #include &lt;ATen/Parallel.h&gt; // Required for get_num_threads and set_num_threads
    #include &lt;random&gt; 
    std::vector&lt;float&gt; batch_infer_distance(torch::jit::Module&amp; model, const std::vector&lt;std::vector&lt;float&gt;&gt;&amp; point_pairs) {
        int batch_size = point_pairs.size();
        torch::Tensor input_tensor = torch::empty({batch_size, 4});
        for (int i = 0; i &lt; batch_size; ++i) {
            input_tensor[i] = torch::tensor(point_pairs[i]);
        }
        
        torch::Tensor output = model.forward({input_tensor}).toTensor();
        std::vector&lt;float&gt; distances(output.data_ptr&lt;float&gt;(), output.data_ptr&lt;float&gt;() + output.numel());
        return distances;
    }


    std::vector&lt;std::vector&lt;float&gt;&gt; generate_random_points(int batch_size) {
        std::vector&lt;std::vector&lt;float&gt;&gt; point_pairs;
        point_pairs.reserve(batch_size);
        std::random_device rd;  
        std::mt19937 gen(rd()); 
        std::uniform_real_distribution&lt;&gt; dis(0.0, 100.0); // Uniform distribution between 0 and 100
        for (int i = 0; i &lt; batch_size; ++i) {
            float x1 = dis(gen);
            float y1 = dis(gen);
            float x2 = dis(gen);
            float y2 = dis(gen);
            point_pairs.push_back({x1, y1, x2, y2});
        }

        return point_pairs;
    }

    int main(int argc, const char* argv[]) {

        int num_threads = at::get_num_threads();
        
        // Print the number of threads
        std::cout &lt;&lt; "Current number of threads: " &lt;&lt; num_threads &lt;&lt; std::endl;

        torch::jit::Module model;
        try {
            model = torch::jit::load("distance_model_traced.pt");
        } catch (const c10::Error&amp; e) {
            std::cerr &lt;&lt; "Error loading the model\n";
            return -1;
        }
        
        std::vector&lt;int&gt; batch_sizes = {1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024}; // Example batch sizes to test
        
        for (int batch_size : batch_sizes) {
            double total_time = 0.0;
            int trial_num = 100000;

            for (int trial = 0; trial &lt; trial_num; ++trial) {
                // Generate random input points for this batch size
                std::vector&lt;std::vector&lt;float&gt;&gt; point_pairs = generate_random_points(batch_size);
                
                // Start timing
                auto start = std::chrono::high_resolution_clock::now();
                
                // Perform inference
                batch_infer_distance(model, point_pairs);
                
                // End timing
                auto end = std::chrono::high_resolution_clock::now();
                std::chrono::duration&lt;double&gt; duration = end - start;
                
                // Accumulate the time for this trial
                total_time += duration.count();
            }
            double average_time = total_time / trial_num;
            
            // Print the average time for the batch size
            std::cout &lt;&lt; "Batch Size: " &lt;&lt; batch_size &lt;&lt; ", Average Time: " &lt;&lt; average_time &lt;&lt; " seconds" &lt;&lt; std::endl;
        }


        return 0;
    }
</code></pre>
<pre><code class="lang-auto">Current number of threads: 4
Model loaded successfully.
Batch Size: 1, Average Time: 1.86423e-05 seconds, Throughput: 53641.4 samples/second
Batch Size: 2, Average Time: 1.90694e-05 seconds, Throughput: 104880 samples/second
Batch Size: 4, Average Time: 2.47976e-05 seconds, Throughput: 161306 samples/second
Batch Size: 8, Average Time: 3.45569e-05 seconds, Throughput: 231502 samples/second
Batch Size: 16, Average Time: 6.15277e-05 seconds, Throughput: 260046 samples/second
Batch Size: 32, Average Time: 9.73409e-05 seconds, Throughput: 328742 samples/second
Batch Size: 64, Average Time: 0.000174497 seconds, Throughput: 366770 samples/second
Batch Size: 128, Average Time: 0.000330429 seconds, Throughput: 387375 samples/second
Batch Size: 256, Average Time: 0.000638286 seconds, Throughput: 401074 samples/second
Batch Size: 512, Average Time: 0.00125685 seconds, Throughput: 407367 samples/second
Batch Size: 1024, Average Time: 0.00267883 seconds, Throughput: 382256 samples/second
</code></pre>
<p>My ultimate goal is to decrease the size of aws server or the number of aws server I am using.</p>
<p>My question is</p>
<ol>
<li>It would seem like in this model given the current feature size , batch size of 512 is the most optimal (since 1ms is acceptable for our api service), am i measuring performance correctly?</li>
<li>What controls throughput in this case? What exactly is making bigger batch obtaining higher throughput and why is batch of 1024 have lower throughput than 512<br>
My guess is that it has to do with the size of the model + size of feature + batch size fitting into l1 cache , and simd , is there other factor?</li>
<li>Is there anything I am missing with performance tuning , I am benchmarking inference only right now, how would I bench mark thrift + inference, I would probably allocate X thread for thrift and Y thread for inference but how do I find the optimal number?</li>
</ol>