Spaces:
Running
Running
File size: 2,117 Bytes
f3270e6 |
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 |
# Copyright (C) 2021-2025, Mindee.
# This program is licensed under the Apache License 2.0.
# See LICENSE or go to <https://opensource.org/licenses/Apache-2.0> for full license details.
"""Text recognition latency benchmark"""
import argparse
import time
import numpy as np
import torch
from doctr.models import recognition
@torch.inference_mode()
def main(args):
device = torch.device("cuda:0" if args.gpu else "cpu")
# Pretrained imagenet model
model = (
recognition.__dict__[args.arch](
pretrained=args.pretrained,
pretrained_backbone=False,
)
.eval()
.to(device=device)
)
# Input
img_tensor = torch.rand((args.batch_size, 3, args.size, 4 * args.size)).to(device=device)
# Warmup
for _ in range(10):
_ = model(img_tensor)
timings = []
# Evaluation runs
for _ in range(args.it):
start_ts = time.perf_counter()
_ = model(img_tensor)
timings.append(time.perf_counter() - start_ts)
_timings = np.array(timings)
print(f"{args.arch} ({args.it} runs on ({args.size}, {4 * args.size}) inputs in batches of {args.batch_size})")
print(f"mean {1000 * _timings.mean():.2f}ms, std {1000 * _timings.std():.2f}ms")
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="docTR latency benchmark for text recognition (PyTorch)",
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)
parser.add_argument("arch", type=str, help="Architecture to use")
parser.add_argument("--batch-size", "-b", type=int, default=64, help="The batch_size")
parser.add_argument("--size", type=int, default=32, help="The image input size")
parser.add_argument("--gpu", dest="gpu", help="Should the benchmark be performed on GPU", action="store_true")
parser.add_argument("--it", type=int, default=100, help="Number of iterations to run")
parser.add_argument(
"--pretrained", dest="pretrained", help="Use pre-trained models from the modelzoo", action="store_true"
)
args = parser.parse_args()
main(args)
|