* Construct a general training worker * @param _solver pointer to graph embedding solver * @param _device_id GPU id */
| 1229 | * @param _device_id GPU id |
| 1230 | */ |
| 1231 | WorkerMixin(Solver *_solver, int _device_id) : |
| 1232 | solver(_solver), device_id(_device_id), negative_sampler(device_id), |
| 1233 | batch(device_id), negative_batch(device_id), logits(device_id), loss(device_id), random(device_id) { |
| 1234 | CUDA_CHECK(cudaSetDevice(device_id)); |
| 1235 | |
| 1236 | CUDA_CHECK(cudaStreamCreate(&work_stream)); |
| 1237 | CUDA_CHECK(cudaStreamCreate(&sample_stream)); |
| 1238 | |
| 1239 | // work stream |
| 1240 | batch.stream = work_stream; |
| 1241 | negative_batch.stream = work_stream; |
| 1242 | logits.stream = work_stream; |
| 1243 | loss.stream = work_stream; |
| 1244 | // sample stream |
| 1245 | negative_sampler.stream = sample_stream; |
| 1246 | random.stream = sample_stream; |
| 1247 | CURAND_CHECK(curandCreateGenerator(&generator, CURAND_RNG_PSEUDO_DEFAULT)); |
| 1248 | std::uniform_int_distribution<unsigned long long> random_seed(0, ULLONG_MAX); |
| 1249 | CURAND_CHECK(curandSetPseudoRandomGeneratorSeed(generator, random_seed(seed))); |
| 1250 | CURAND_CHECK(curandSetStream(generator, sample_stream)); |
| 1251 | } |
| 1252 | |
| 1253 | WorkerMixin(const WorkerMixin &) = delete; |
| 1254 |
nothing calls this directly
no outgoing calls
no test coverage detected