forked from timesler/facenet-pytorch
-
Notifications
You must be signed in to change notification settings - Fork 0
/
perf_test.py
39 lines (31 loc) · 966 Bytes
/
perf_test.py
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
from facenet_pytorch import MTCNN, training
import torch
from torchvision import datasets, transforms
from torch.utils.data import DataLoader, RandomSampler
from tqdm import tqdm
import time
def main():
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f'Running on device "{device}"')
mtcnn = MTCNN(device=device)
batch_size = 32
# Generate data loader
ds = datasets.ImageFolder(
root='data/test_images/',
transform=transforms.Resize((512, 512))
)
dl = DataLoader(
dataset=ds,
num_workers=4,
collate_fn=training.collate_pil,
batch_size=batch_size,
sampler=RandomSampler(ds, replacement=True, num_samples=960),
)
start = time.time()
faces = []
for x, _ in tqdm(dl):
faces.extend(mtcnn(x))
elapsed = time.time() - start
print(f'Elapsed: {elapsed} | EPS: {len(dl) * batch_size / elapsed}')
if __name__ == '__main__':
main()