()
| 22 | |
| 23 | |
| 24 | def main(): |
| 25 | args = parse_args() |
| 26 | |
| 27 | cfg = Config.fromfile(args.config) |
| 28 | # set cudnn_benchmark |
| 29 | torch.backends.cudnn.benchmark = False |
| 30 | cfg.model.pretrained = None |
| 31 | cfg.data.test.test_mode = True |
| 32 | |
| 33 | # build the dataloader |
| 34 | # TODO: support multiple images per gpu (only minor changes are needed) |
| 35 | dataset = build_dataset(cfg.data.test) |
| 36 | data_loader = build_dataloader( |
| 37 | dataset, |
| 38 | samples_per_gpu=1, |
| 39 | workers_per_gpu=cfg.data.workers_per_gpu, |
| 40 | dist=False, |
| 41 | shuffle=False) |
| 42 | |
| 43 | # build the model and load checkpoint |
| 44 | cfg.model.train_cfg = None |
| 45 | model = build_segmentor(cfg.model, test_cfg=cfg.get('test_cfg')) |
| 46 | fp16_cfg = cfg.get('fp16', None) |
| 47 | if fp16_cfg is not None: |
| 48 | wrap_fp16_model(model) |
| 49 | load_checkpoint(model, args.checkpoint, map_location='cpu') |
| 50 | |
| 51 | model = MMDataParallel(model, device_ids=[0]) |
| 52 | |
| 53 | model.eval() |
| 54 | |
| 55 | # the first several iterations may be very slow so skip them |
| 56 | num_warmup = 5 |
| 57 | pure_inf_time = 0 |
| 58 | total_iters = 200 |
| 59 | |
| 60 | # benchmark with 200 image and take the average |
| 61 | for i, data in enumerate(data_loader): |
| 62 | |
| 63 | torch.cuda.synchronize() |
| 64 | start_time = time.perf_counter() |
| 65 | |
| 66 | with torch.no_grad(): |
| 67 | model(return_loss=False, rescale=True, **data) |
| 68 | |
| 69 | torch.cuda.synchronize() |
| 70 | elapsed = time.perf_counter() - start_time |
| 71 | |
| 72 | if i >= num_warmup: |
| 73 | pure_inf_time += elapsed |
| 74 | if (i + 1) % args.log_interval == 0: |
| 75 | fps = (i + 1 - num_warmup) / pure_inf_time |
| 76 | print(f'Done image [{i + 1:<3}/ {total_iters}], ' |
| 77 | f'fps: {fps:.2f} img / s') |
| 78 | |
| 79 | if (i + 1) == total_iters: |
| 80 | fps = (i + 1 - num_warmup) / pure_inf_time |
| 81 | print(f'Overall fps: {fps:.2f} img / s') |
no test coverage detected