(use_cuda)
| 21 | |
| 22 | |
| 23 | def test(use_cuda): |
| 24 | torch.set_grad_enabled(False) |
| 25 | torch.set_num_threads(4) |
| 26 | turbo_transformers.set_num_threads(4) |
| 27 | |
| 28 | test_device = torch.device('cuda:0') if use_cuda else \ |
| 29 | torch.device('cpu:0') |
| 30 | |
| 31 | cfg = RobertaConfig() |
| 32 | torch_model = RobertaModel(cfg) |
| 33 | torch_model.eval() |
| 34 | |
| 35 | if torch.cuda.is_available(): |
| 36 | torch_model.to(test_device) |
| 37 | |
| 38 | turbo_model = turbo_transformers.RobertaModel.from_torch( |
| 39 | torch_model, test_device) |
| 40 | |
| 41 | input_ids = torch.randint(low=0, |
| 42 | high=cfg.vocab_size - 1, |
| 43 | size=(1, 10), |
| 44 | dtype=torch.long, |
| 45 | device=test_device) |
| 46 | |
| 47 | torch_result = torch_model(input_ids) |
| 48 | torch_result_final = torch_result[0].cpu().numpy() |
| 49 | |
| 50 | turbo_result = turbo_model(input_ids) |
| 51 | turbo_result_final = turbo_result[0].cpu().numpy() |
| 52 | |
| 53 | # See the differences |
| 54 | # print(numpy.size(torch_result_final), numpy.size(turbo_result_final)) |
| 55 | print(torch_result_final - turbo_result_final) |
| 56 | assert (numpy.allclose(torch_result_final, |
| 57 | turbo_result_final, |
| 58 | atol=1e-3, |
| 59 | rtol=1e-3)) |
| 60 | |
| 61 | |
| 62 | if __name__ == "__main__": |
no test coverage detected