| 53 | } |
| 54 | |
| 55 | bool test_multiple_threads(const std::string &model_path, bool only_input, |
| 56 | bool use_cuda, int n_threads) { |
| 57 | std::shared_ptr<BertModel> model_ptr = |
| 58 | std::make_shared<BertModel>(model_path, DLDeviceType::kDLCPU, 12, 12); |
| 59 | // input_ids, position_ids, segment_ids lengths of each row may not be the |
| 60 | // same. For example. std::vector<std::vector<int64_t>> input_ids{{1, 2, 3, 4, |
| 61 | // 5, 6, 7}, |
| 62 | // {1, 2}}; |
| 63 | std::vector<std::vector<int64_t>> input_ids{{12166, 10699, 16752, 4454}, |
| 64 | {5342, 16471, 817, 16022}}; |
| 65 | std::vector<std::vector<int64_t>> position_ids{{1, 0, 0, 0}, {1, 1, 1, 0}}; |
| 66 | std::vector<std::vector<int64_t>> segment_ids{{1, 1, 1, 0}, {1, 0, 0, 0}}; |
| 67 | if (only_input) { |
| 68 | position_ids.clear(); |
| 69 | segment_ids.clear(); |
| 70 | } |
| 71 | std::vector<std::thread> threads; |
| 72 | threads.reserve(n_threads); |
| 73 | |
| 74 | std::vector<std::future<std::vector<float>>> result_list; |
| 75 | for (int i = 0; i < n_threads; ++i) { |
| 76 | std::packaged_task<std::vector<float>( |
| 77 | const std::shared_ptr<BertModel>, |
| 78 | const std::vector<std::vector<int64_t>> &, |
| 79 | const std::vector<std::vector<int64_t>> &, |
| 80 | const std::vector<std::vector<int64_t>> &, PoolType, bool)> |
| 81 | task(CallBackFunction); |
| 82 | result_list.emplace_back(task.get_future()); |
| 83 | threads.emplace_back(std::thread(std::move(task), model_ptr, input_ids, |
| 84 | position_ids, segment_ids, |
| 85 | PoolType::kFirst, true)); |
| 86 | } |
| 87 | |
| 88 | for (int i = 0; i < n_threads; ++i) { |
| 89 | auto vec = result_list[i].get(); |
| 90 | assert(vec.size() == 768 * 2); |
| 91 | |
| 92 | for (size_t i = 0; i < vec.size(); ++i) { |
| 93 | assert(!std::isnan(vec.data()[i])); |
| 94 | assert(!std::isinf(vec.data()[i])); |
| 95 | } |
| 96 | // Attention, the hard code value is based on huggingface/transformers |
| 97 | // bert-base-uncased (2020.04.23 version), you may need to change it to |
| 98 | // real-time values. |
| 99 | if (only_input) { |
| 100 | std::cerr << vec.data()[0] << std::endl; |
| 101 | std::cerr << vec.data()[1] << std::endl; |
| 102 | std::cerr << vec.data()[768] << std::endl; |
| 103 | std::cerr << vec.data()[768 + 1] << std::endl; |
| 104 | assert(fabs(vec.data()[0] - -0.1901) < 1e-3); |
| 105 | assert(fabs(vec.data()[1] - 0.0193) < 1e-3); |
| 106 | assert(fabs(vec.data()[768] - 0.3060) < 1e-3); |
| 107 | assert(fabs(vec.data()[768 + 1] - 0.1162) < 1e-3); |
| 108 | } else { |
| 109 | assert(fabs(vec.data()[0] - -0.5503) < 1e-3); |
| 110 | assert(fabs(vec.data()[1] - 0.1295) < 1e-3); |
| 111 | assert(fabs(vec.data()[768] - -0.5545) < 1e-3); |
| 112 | assert(fabs(vec.data()[768 + 1] - -0.1182) < 1e-3); |