python用装饰器统计函数运行时间

time_count_via_decorator.py

import time

def clock(func):
    def clocked(*args, **kw):
        t0 = time.perf_counter()
        result = func(*args, **kw)
        elapsed = time.perf_counter() - t0
        name = func.__name__
        arg_str = ', '.join(repr(arg) for arg in args)
        print('[%0.8fs] %s' % (elapsed, name))
        return result
    return clocked

@clock
def main(model, X):
    out = model(X)
    print(out.shape)
    

device = torch.device("cuda:0")
model = nn.Linear(32, 12).to(device)
X = torch.randn(207, 32).to(device)
main(model, X)         # model, X is the args will passed to clocked
# > output: [0.97528419s] main