(loadtype: LoadType, use_cuda: bool)
| 24 | |
| 25 | |
| 26 | def test(loadtype: LoadType, use_cuda: bool): |
| 27 | model_id = "bert-base-uncased" |
| 28 | model = transformers.BertModel.from_pretrained(model_id) |
| 29 | model.eval() |
| 30 | torch.set_grad_enabled(False) |
| 31 | |
| 32 | test_device = torch.device('cuda:0') if use_cuda else \ |
| 33 | torch.device('cpu:0') |
| 34 | |
| 35 | cfg = model.config |
| 36 | # use 4 threads for computing |
| 37 | turbo_transformers.set_num_threads(4) |
| 38 | |
| 39 | input_ids = torch.tensor( |
| 40 | ([12166, 10699, 16752, 4454], [5342, 16471, 817, 16022]), |
| 41 | dtype=torch.long) |
| 42 | # position_ids = torch.tensor(([1, 0, 0, 0], [1, 1, 1, 0]), dtype=torch.long) |
| 43 | segment_ids = torch.tensor(([1, 1, 1, 0], [1, 0, 0, 0]), dtype=torch.long) |
| 44 | |
| 45 | start_time = time.time() |
| 46 | for _ in range(10): |
| 47 | torch_res = model( |
| 48 | input_ids, token_type_ids=segment_ids |
| 49 | ) # sequence_output, pooled_output, (hidden_states), (attentions) |
| 50 | end_time = time.time() |
| 51 | print("\ntorch time consum: {}".format(end_time - start_time)) |
| 52 | print("torch bert sequence output: ", |
| 53 | torch_res[0][:, 0, :]) #get the first sequence |
| 54 | print("torch bert pooler output: ", torch_res[1]) # pooled_output |
| 55 | |
| 56 | # there are three ways to load pretrained model. |
| 57 | if loadtype is LoadType.PYTORCH: |
| 58 | # 1, from a PyTorch model, which has loaded a pretrained model |
| 59 | # note that you can choose "turbo" or "onnxrt" as backend |
| 60 | # "turbo" is a hand-crafted implementation and optimized with OMP. |
| 61 | tt_model = turbo_transformers.BertModel.from_torch( |
| 62 | model, test_device, "onnxrt") |
| 63 | elif loadtype is LoadType.PRETRAINED: |
| 64 | # 2. directly load from checkpoint (torch saved model) |
| 65 | tt_model = turbo_transformers.BertModel.from_pretrained( |
| 66 | model_id, test_device) |
| 67 | elif loadtype is LoadType.NPZ: |
| 68 | # 3. load model from npz |
| 69 | if len(sys.argv) == 2: |
| 70 | try: |
| 71 | print(sys.argv[1]) |
| 72 | in_file = sys.argv[1] |
| 73 | except: |
| 74 | sys.exit("ERROR. can not open ", sys.argv[1]) |
| 75 | else: |
| 76 | in_file = "/workspace/bert_torch.npz" |
| 77 | tt_model = turbo_transformers.BertModel.from_npz( |
| 78 | in_file, cfg, test_device) |
| 79 | else: |
| 80 | raise ("LoadType is not supported") |
| 81 | |
| 82 | start_time = time.time() |
| 83 | for _ in range(10): |
no test coverage detected