KieDani/SegformerPlusPlus

1from typing import Union, List, Tuple
2
3import numpy as np
4import torch
5
6from .utils import benchmark
7
8device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
9
10
11def random_benchmark(
12        model: torch.nn.Module,
13        batch_size: Union[int, List[int]] = 1,
14        image_size: Union[Tuple[int], List[Tuple[int]]] = (3, 1024, 1024),
15):
16    """
17    Calculate the FPS of a given model using randomly generated tensors.
18
19    Args:
20        model: instance of a model (e.g. SegFormer)
21        batch_size: the batch size(s) at which to calculate the FPS (e.g. 1 or [1, 2, 4])
22        image_size: the size of the images to use (e.g. (3, 1024, 1024))
23
24    Returns: the FPS values calculated for all image sizes and batch sizes in the form of a dictionary
25
26    """
27    if isinstance(batch_size, int):
28        batch_size = [batch_size]
29    if isinstance(image_size, tuple):
30        image_size = [image_size]
31
32    values = {}
33    throughput_values = []
34
35    for i in image_size:
36        # fill with fps for each batch size
37        fps = []
38        for b in batch_size:
39            for _ in range(4):
40                # Baseline benchmark
41                if i[1] >= 1024:
42                    r = 16
43                else:
44                    r = 32
45                baseline_throughput = benchmark(
46                    model.to(device),
47                    device=device,
48                    verbose=True,
49                    runs=r,
50                    batch_size=b,
51                    input_size=i
52                )
53                throughput_values.append(baseline_throughput)
54            throughput_values = np.asarray(throughput_values)
55            throughput = np.around(np.mean(throughput_values), decimals=2)
56            print('Im_size:', i, 'Batch_size:', b, 'Mean:', throughput, 'Std:',
57                  np.around(np.std(throughput_values), decimals=2))
58            throughput_values = []
59            fps.append({b: throughput})
60        values[i] = fps
61    return values
62