File size: 3,660 Bytes
e98bd8c 85ebba9 e98bd8c 85ebba9 e98bd8c 10e4799 e98bd8c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 |
import torch
from PIL import Image
import torchvision.transforms as T
import os
from typing import Union, List, Tuple
import numpy as np
from .utils.benchmark import benchmark
# Gerät auswählen
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {device}")
if device.type == 'cuda':
print(f"CUDA Device Name: {torch.cuda.get_device_name(torch.cuda.current_device())}")
torch.manual_seed(42)
torch.cuda.manual_seed_all(42)
def cityscape_benchmark(
model: torch.nn.Module,
image_path: str,
batch_size: Union[int, List[int]] = 1,
image_size: Union[Tuple[int], List[Tuple[int]]] = (3, 1024, 1024),
save_output: bool = True,
):
"""
Calculate the FPS of a given model using an actual Cityscapes image.
Args:
model: instance of a model (e.g. SegFormer)
image_path: the path to the Cityscapes image
batch_size: the batch size(s) at which to calculate the FPS (e.g. 1 or [1, 2, 4])
image_size: the size of the images to use (e.g. (3, 1024, 1024))
save_output: whether to save the output prediction (default True)
Returns:
the FPS values calculated for all image sizes and batch sizes in the form of a dictionary
"""
if isinstance(batch_size, int):
batch_size = [batch_size]
if isinstance(image_size, tuple):
image_size = [image_size]
values = {}
throughput_values = []
model = model.to(device)
model.eval()
assert os.path.exists(image_path), f"Image not found: {image_path}"
image = Image.open(image_path).convert("RGB")
img_tensor = T.ToTensor()(image)
mean = img_tensor.mean(dim=(1, 2))
std = img_tensor.std(dim=(1, 2))
print(f"Calculated Mean: {mean}")
print(f"Calculated Std: {std}")
transform = T.Compose([
T.Resize((image_size[0][1], image_size[0][2])),
T.ToTensor(),
T.Normalize(mean=mean.tolist(),
std=std.tolist())
])
img_tensor = transform(image).unsqueeze(0).to(device)
for i in image_size:
# fill with fps for each batch size
fps = []
for b in batch_size:
for _ in range(4):
# Baseline benchmark
if i[1] >= 1024:
r = 16
else:
r = 32
baseline_throughput = benchmark(
model.to(device),
device=device,
verbose=True,
runs=r,
batch_size=b,
input_size=i
)
throughput_values.append(baseline_throughput)
throughput_values = np.asarray(throughput_values)
throughput = np.around(np.mean(throughput_values), decimals=2)
print('Im_size:', i, 'Batch_size:', b, 'Mean:', throughput, 'Std:',
np.around(np.std(throughput_values), decimals=2))
throughput_values = []
fps.append({b: throughput})
values[i] = fps
if save_output:
with torch.no_grad():
output = model(img_tensor)
pred = torch.argmax(output, dim=1).squeeze(0).cpu().numpy()
# Speichere Prediction als Text ab
cwd = os.getcwd()
output_path = os.path.join(cwd, 'segformer_plusplus', 'cityscapes_prediction_output.txt')
np.savetxt(output_path, pred, fmt="%d")
print("Prediction saved as cityscapes_prediction_output.txt")
return values
|