HairFastGAN

1
import functools
2
import sys
3
import time
4

5
import numpy as np
6
import torch
7

8

9
def get_time():
10
    torch.cuda.current_stream().synchronize()
11
    return time.time()
12

13

14
def bench_session(func):
15
    times = []
16

17
    @functools.wraps(func)
18
    def wraps(*args, **kwargs):
19
        if kwargs.pop('benchmark', False):
20
            nonlocal times
21
            start = get_time()
22

23
            result = func(*args, **kwargs)
24

25
            eval_time = get_time() - start
26
            times.append(eval_time)
27

28
            print(f'\n{len(times)} experiment ended in {eval_time:.3f}(s)', file=sys.stderr)
29
            print(f'min time: {np.min(times):.3f}(s),'
30
                  f' median time: {np.median(times):.3f}(s),'
31
                  f' std time: {np.std(times):.3f}(s)', file=sys.stderr)
32
            return result
33
        else:
34
            return func(*args, **kwargs)
35

36
    return wraps
37