⚡️ Speed up method ClassifierHead.forward by 7%

[codeflash-ai]

📄 7% (0.07x) speedup for ClassifierHead.forward in doctr/models/classification/vit/pytorch.py

⏱️ Runtime : 2.64 milliseconds → 2.47 milliseconds (best of 44 runs)

📝 Explanation and details

The optimization replaces PyTorch tensor slicing x[:, 0] with the more efficient x.select(1, 0) method to extract the first token along dimension 1.

Key optimization:

• x.select(1, 0) is a direct indexing operation that operates at the C++ backend level
• x[:, 0] creates an intermediate view through Python's slicing mechanism before extracting the data
• The select method bypasses the overhead of Python slice object creation and view management

Why it's faster:
The line profiler shows the slicing operation (x[:, 0]) took 116,042 ns per hit, while x.select(1, 0) takes only 18,885 ns per hit - a 6x reduction in per-operation cost. This translates to the overall 7% speedup.

Performance characteristics from tests:

• Consistent 10-15% improvements across most test cases
• Best performance gains (20-25%) on smaller tensors and edge cases like minimal inputs and empty batches
• Even large-scale tests (100MB tensors) show measurable improvements (1-2%)
• The optimization maintains identical behavior and error handling

This is particularly beneficial for Vision Transformer classification heads where this operation runs frequently during inference, as it extracts the classification token (first position) from the sequence for final prediction.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 71 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Generated Regression Tests and Runtime

import pytest  # used for our unit tests
import torch
from doctr.models.classification.vit.pytorch import ClassifierHead
from torch import nn

# unit tests

# ---- BASIC TEST CASES ----

def test_forward_basic_shape_and_type():
    # Test that output shape is correct for a standard input
    batch_size = 4
    seq_len = 10
    in_channels = 8
    num_classes = 3
    model = ClassifierHead(in_channels, num_classes)
    # Create random input tensor
    x = torch.randn(batch_size, seq_len, in_channels)
    codeflash_output = model.forward(x); out = codeflash_output # 55.3μs -> 48.7μs (13.6% faster)

def test_forward_basic_values_consistency():
    # Test that output for repeated input is consistent
    batch_size = 2
    seq_len = 5
    in_channels = 4
    num_classes = 2
    model = ClassifierHead(in_channels, num_classes)
    # All inputs are identical
    x = torch.ones(batch_size, seq_len, in_channels)
    codeflash_output = model.forward(x); out = codeflash_output # 46.5μs -> 41.3μs (12.7% faster)

def test_forward_basic_single_batch():
    # Test with batch_size=1
    batch_size = 1
    seq_len = 6
    in_channels = 7
    num_classes = 5
    model = ClassifierHead(in_channels, num_classes)
    x = torch.randn(batch_size, seq_len, in_channels)
    codeflash_output = model.forward(x); out = codeflash_output # 42.8μs -> 36.6μs (17.1% faster)

def test_forward_basic_single_class():
    # Test with num_classes=1
    batch_size = 3
    seq_len = 8
    in_channels = 6
    num_classes = 1
    model = ClassifierHead(in_channels, num_classes)
    x = torch.randn(batch_size, seq_len, in_channels)
    codeflash_output = model.forward(x); out = codeflash_output # 39.5μs -> 34.9μs (13.3% faster)

# ---- EDGE TEST CASES ----

def test_forward_edge_minimal_input():
    # Test with batch_size=1, seq_len=1, in_channels=1, num_classes=1
    model = ClassifierHead(1, 1)
    x = torch.tensor([[[2.0]]])  # shape (1, 1, 1)
    codeflash_output = model.forward(x); out = codeflash_output # 38.2μs -> 33.9μs (12.6% faster)

def test_forward_edge_zero_batch():
    # Test with batch_size=0
    batch_size = 0
    seq_len = 5
    in_channels = 3
    num_classes = 2
    model = ClassifierHead(in_channels, num_classes)
    x = torch.randn(batch_size, seq_len, in_channels)
    codeflash_output = model.forward(x); out = codeflash_output # 33.9μs -> 27.2μs (24.5% faster)

def test_forward_edge_zero_seq_len():
    # Test with seq_len=0 (should raise IndexError)
    batch_size = 2
    seq_len = 0
    in_channels = 4
    num_classes = 3
    model = ClassifierHead(in_channels, num_classes)
    x = torch.randn(batch_size, seq_len, in_channels)
    # Should raise IndexError because x[:, 0] is out of bounds
    with pytest.raises(IndexError):
        model.forward(x) # 52.9μs -> 61.5μs (13.9% slower)

def test_forward_edge_wrong_input_shape():
    # Test with input shape missing in_channels dimension
    batch_size = 3
    seq_len = 5
    in_channels = 4
    num_classes = 2
    model = ClassifierHead(in_channels, num_classes)
    # Incorrect shape: missing in_channels
    x = torch.randn(batch_size, seq_len)
    with pytest.raises(RuntimeError):
        model.forward(x) # 104μs -> 97.4μs (6.89% faster)

def test_forward_edge_incorrect_in_channels():
    # Test with input in_channels not matching model in_channels
    batch_size = 2
    seq_len = 3
    in_channels = 5
    num_classes = 2
    model = ClassifierHead(in_channels, num_classes)
    # Incorrect in_channels
    x = torch.randn(batch_size, seq_len, in_channels + 1)
    with pytest.raises(RuntimeError):
        model.forward(x) # 91.5μs -> 84.7μs (8.06% faster)

def test_forward_edge_negative_input():
    # Test with negative values in input
    batch_size = 2
    seq_len = 4
    in_channels = 3
    num_classes = 2
    model = ClassifierHead(in_channels, num_classes)
    x = -torch.abs(torch.randn(batch_size, seq_len, in_channels))
    codeflash_output = model.forward(x); out = codeflash_output # 49.3μs -> 43.7μs (13.0% faster)

def test_forward_edge_large_values():
    # Test with very large values in input
    batch_size = 2
    seq_len = 3
    in_channels = 4
    num_classes = 2
    model = ClassifierHead(in_channels, num_classes)
    x = torch.full((batch_size, seq_len, in_channels), 1e6)
    codeflash_output = model.forward(x); out = codeflash_output # 45.8μs -> 37.9μs (20.7% faster)


def test_forward_edge_dtype_half():
    # Test with half precision input
    batch_size = 2
    seq_len = 3
    in_channels = 4
    num_classes = 2
    model = ClassifierHead(in_channels, num_classes)
    x = torch.randn(batch_size, seq_len, in_channels, dtype=torch.float16)
    # nn.Linear does not support float16 on CPU, so should raise error
    with pytest.raises(RuntimeError):
        model.forward(x) # 99.4μs -> 93.7μs (6.11% faster)

def test_forward_edge_non_contiguous_input():
    # Test with non-contiguous input tensor
    batch_size = 2
    seq_len = 3
    in_channels = 4
    num_classes = 2
    model = ClassifierHead(in_channels, num_classes)
    x = torch.randn(batch_size, seq_len, in_channels).transpose(0, 1)
    # Should still work
    codeflash_output = model.forward(x.transpose(0, 1)); out = codeflash_output # 48.0μs -> 42.2μs (13.7% faster)

# ---- LARGE SCALE TEST CASES ----

def test_forward_large_batch():
    # Test with large batch size
    batch_size = 1000
    seq_len = 5
    in_channels = 8
    num_classes = 3
    model = ClassifierHead(in_channels, num_classes)
    x = torch.randn(batch_size, seq_len, in_channels)
    codeflash_output = model.forward(x); out = codeflash_output # 60.4μs -> 55.0μs (9.78% faster)

def test_forward_large_seq_len():
    # Test with large sequence length
    batch_size = 10
    seq_len = 1000
    in_channels = 8
    num_classes = 3
    model = ClassifierHead(in_channels, num_classes)
    x = torch.randn(batch_size, seq_len, in_channels)
    codeflash_output = model.forward(x); out = codeflash_output # 50.0μs -> 44.2μs (13.2% faster)

def test_forward_large_in_channels():
    # Test with large number of input channels
    batch_size = 5
    seq_len = 10
    in_channels = 1000
    num_classes = 7
    model = ClassifierHead(in_channels, num_classes)
    x = torch.randn(batch_size, seq_len, in_channels)
    codeflash_output = model.forward(x); out = codeflash_output # 49.0μs -> 44.1μs (11.3% faster)

def test_forward_large_num_classes():
    # Test with large number of output classes
    batch_size = 8
    seq_len = 6
    in_channels = 12
    num_classes = 1000
    model = ClassifierHead(in_channels, num_classes)
    x = torch.randn(batch_size, seq_len, in_channels)
    codeflash_output = model.forward(x); out = codeflash_output # 56.4μs -> 49.9μs (13.0% faster)

def test_forward_large_total_size():
    # Test with tensor close to 100MB
    batch_size = 100
    seq_len = 100
    in_channels = 100
    num_classes = 10
    model = ClassifierHead(in_channels, num_classes)
    x = torch.randn(batch_size, seq_len, in_channels)
    codeflash_output = model.forward(x); out = codeflash_output # 77.0μs -> 69.5μs (10.7% faster)

def test_forward_large_consistency():
    # Test that output for repeated input is consistent for large batch
    batch_size = 500
    seq_len = 20
    in_channels = 10
    num_classes = 5
    model = ClassifierHead(in_channels, num_classes)
    x = torch.ones(batch_size, seq_len, in_channels)
    codeflash_output = model.forward(x); out = codeflash_output # 57.2μs -> 51.9μs (10.2% faster)
    # All outputs should be identical
    for i in range(1, batch_size):
        pass

# ---- ADDITIONAL EDGE CASES ----

def test_forward_empty_tensor():
    # Test with completely empty tensor (should raise IndexError)
    model = ClassifierHead(1, 1)
    x = torch.empty((0, 0, 0))
    with pytest.raises(IndexError):
        model.forward(x) # 53.8μs -> 62.3μs (13.6% slower)

def test_forward_input_requires_grad():
    # Test with input that requires grad
    batch_size = 2
    seq_len = 3
    in_channels = 4
    num_classes = 2
    model = ClassifierHead(in_channels, num_classes)
    x = torch.randn(batch_size, seq_len, in_channels, requires_grad=True)
    codeflash_output = model.forward(x); out = codeflash_output # 53.4μs -> 46.7μs (14.3% faster)


#------------------------------------------------
import pytest  # used for our unit tests
import torch
from doctr.models.classification.vit.pytorch import ClassifierHead
from torch import nn

# unit tests

# ---- Basic Test Cases ----

def test_forward_basic_shape_and_output():
    # Test with batch_size=2, seq_len=4, in_channels=3, num_classes=5
    in_channels = 3
    num_classes = 5
    batch_size = 2
    seq_len = 4
    # Create input tensor
    x = torch.ones(batch_size, seq_len, in_channels)
    # Initialize ClassifierHead
    head = ClassifierHead(in_channels, num_classes)
    # Forward pass
    codeflash_output = head.forward(x); out = codeflash_output # 50.4μs -> 45.0μs (11.8% faster)


def test_forward_basic_different_values():
    # Test with non-uniform input values
    in_channels = 4
    num_classes = 2
    batch_size = 3
    seq_len = 6
    x = torch.arange(batch_size * seq_len * in_channels, dtype=torch.float32).reshape(batch_size, seq_len, in_channels)
    head = ClassifierHead(in_channels, num_classes)
    codeflash_output = head.forward(x); out = codeflash_output # 47.6μs -> 41.9μs (13.6% faster)
    # Check that the output depends only on x[:, 0]
    manual = head.head(x[:, 0])


def test_forward_basic_single_batch():
    # Test with batch_size=1
    in_channels = 2
    num_classes = 3
    batch_size = 1
    seq_len = 5
    x = torch.zeros(batch_size, seq_len, in_channels)
    head = ClassifierHead(in_channels, num_classes)
    codeflash_output = head.forward(x); out = codeflash_output # 37.9μs -> 34.9μs (8.67% faster)
    # Output should be zeros if weights are zeros and bias is zeros
    head.head.weight.data.zero_()
    head.head.bias.data.zero_()
    codeflash_output = head.forward(x); out = codeflash_output # 14.2μs -> 12.5μs (13.8% faster)


# ---- Edge Test Cases ----

def test_forward_edge_minimal_input():
    # Test with minimal valid input: batch_size=1, seq_len=1, in_channels=1, num_classes=1
    in_channels = 1
    num_classes = 1
    batch_size = 1
    seq_len = 1
    x = torch.tensor([[[42.0]]])
    head = ClassifierHead(in_channels, num_classes)
    codeflash_output = head.forward(x); out = codeflash_output # 38.0μs -> 30.9μs (22.7% faster)


def test_forward_edge_seq_len_1():
    # Test with seq_len=1, should only use x[:, 0]
    in_channels = 3
    num_classes = 2
    batch_size = 4
    seq_len = 1
    x = torch.randn(batch_size, seq_len, in_channels)
    head = ClassifierHead(in_channels, num_classes)
    codeflash_output = head.forward(x); out = codeflash_output # 43.1μs -> 37.8μs (14.1% faster)
    # Should be equal to head.head(x.squeeze(1))
    manual = head.head(x.squeeze(1))


def test_forward_edge_in_channels_1():
    # Test with in_channels=1
    in_channels = 1
    num_classes = 3
    batch_size = 2
    seq_len = 5
    x = torch.randn(batch_size, seq_len, in_channels)
    head = ClassifierHead(in_channels, num_classes)
    codeflash_output = head.forward(x); out = codeflash_output # 40.4μs -> 34.1μs (18.2% faster)


def test_forward_edge_num_classes_1():
    # Test with num_classes=1
    in_channels = 4
    num_classes = 1
    batch_size = 2
    seq_len = 3
    x = torch.randn(batch_size, seq_len, in_channels)
    head = ClassifierHead(in_channels, num_classes)
    codeflash_output = head.forward(x); out = codeflash_output # 39.9μs -> 35.2μs (13.4% faster)


def test_forward_edge_empty_batch():
    # Test with batch_size=0
    in_channels = 3
    num_classes = 2
    batch_size = 0
    seq_len = 4
    x = torch.randn(batch_size, seq_len, in_channels)
    head = ClassifierHead(in_channels, num_classes)
    codeflash_output = head.forward(x); out = codeflash_output # 30.6μs -> 26.8μs (14.4% faster)



def test_forward_edge_in_channels_mismatch():
    # Test with input in_channels not matching head.in_channels
    in_channels = 3
    num_classes = 2
    batch_size = 2
    seq_len = 4
    # x has in_channels=4 instead of 3
    x = torch.randn(batch_size, seq_len, 4)
    head = ClassifierHead(in_channels, num_classes)
    # Should raise RuntimeError due to shape mismatch in Linear
    with pytest.raises(RuntimeError):
        head.forward(x) # 101μs -> 95.6μs (5.83% faster)


# ---- Large Scale Test Cases ----

def test_forward_large_batch_size():
    # Test with large batch size
    in_channels = 8
    num_classes = 10
    batch_size = 1000  # max allowed
    seq_len = 5
    x = torch.randn(batch_size, seq_len, in_channels)
    head = ClassifierHead(in_channels, num_classes)
    codeflash_output = head.forward(x); out = codeflash_output # 65.3μs -> 60.0μs (8.90% faster)


def test_forward_large_seq_len():
    # Test with large sequence length
    in_channels = 16
    num_classes = 4
    batch_size = 10
    seq_len = 1000  # max allowed
    x = torch.randn(batch_size, seq_len, in_channels)
    head = ClassifierHead(in_channels, num_classes)
    codeflash_output = head.forward(x); out = codeflash_output # 51.8μs -> 45.9μs (12.8% faster)


def test_forward_large_in_channels():
    # Test with large in_channels
    in_channels = 512
    num_classes = 32
    batch_size = 4
    seq_len = 3
    x = torch.randn(batch_size, seq_len, in_channels)
    head = ClassifierHead(in_channels, num_classes)
    codeflash_output = head.forward(x); out = codeflash_output # 50.4μs -> 44.4μs (13.4% faster)


def test_forward_large_num_classes():
    # Test with large num_classes
    in_channels = 32
    num_classes = 512
    batch_size = 2
    seq_len = 7
    x = torch.randn(batch_size, seq_len, in_channels)
    head = ClassifierHead(in_channels, num_classes)
    codeflash_output = head.forward(x); out = codeflash_output # 52.8μs -> 46.7μs (12.9% faster)


def test_forward_large_tensor_memory_limit():
    # Ensure tensor size does not exceed 100MB
    # Each float32 is 4 bytes, so max elements = 100MB / 4 = 25_000_000
    # Let's use batch_size=100, seq_len=10, in_channels=25000 (100*10*25000=25_000_000)
    in_channels = 25000
    num_classes = 2
    batch_size = 100
    seq_len = 10
    x = torch.randn(batch_size, seq_len, in_channels)
    head = ClassifierHead(in_channels, num_classes)
    codeflash_output = head.forward(x); out = codeflash_output # 776μs -> 767μs (1.26% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-ClassifierHead.forward-mg7qwcn6 and push.