(loadtype: LoadType, use_cuda: bool)
| 25 | |
| 26 | |
| 27 | def test(loadtype: LoadType, use_cuda: bool): |
| 28 | cfg = transformers.AlbertConfig() |
| 29 | model = transformers.AlbertModel(cfg) |
| 30 | model.eval() |
| 31 | torch.set_grad_enabled(False) |
| 32 | |
| 33 | test_device = torch.device('cuda:0') if use_cuda else \ |
| 34 | torch.device('cpu:0') |
| 35 | |
| 36 | cfg = model.config |
| 37 | # use 4 threads for computing |
| 38 | turbo_transformers.set_num_threads(4) |
| 39 | |
| 40 | input_ids = torch.tensor( |
| 41 | ([12166, 10699, 16752, 4454], [5342, 16471, 817, 16022]), |
| 42 | dtype=torch.long) |
| 43 | model.to(test_device) |
| 44 | start_time = time.time() |
| 45 | for _ in range(10): |
| 46 | torch_res = model(input_ids) |
| 47 | end_time = time.time() |
| 48 | print("\ntorch time consum: {}".format(end_time - start_time)) |
| 49 | |
| 50 | # there are three ways to load pretrained model. |
| 51 | if loadtype is LoadType.PYTORCH: |
| 52 | # 1, from a PyTorch model, which has loaded a pretrained model |
| 53 | tt_model = turbo_transformers.AlbertModel.from_torch(model) |
| 54 | else: |
| 55 | raise ("LoadType is not supported") |
| 56 | |
| 57 | start_time = time.time() |
| 58 | for _ in range(10): |
| 59 | res = tt_model(input_ids) # sequence_output, pooled_output |
| 60 | end_time = time.time() |
| 61 | |
| 62 | print("\nturbo time consum: {}".format(end_time - start_time)) |
| 63 | assert (numpy.max( |
| 64 | numpy.abs(res[0].cpu().numpy() - torch_res[0].cpu().numpy())) < 0.1) |
| 65 | |
| 66 | |
| 67 | if __name__ == "__main__": |
no test coverage detected