Determine and allocate all resources for the worker */
| 1279 | |
| 1280 | /** Determine and allocate all resources for the worker */ |
| 1281 | void build() { |
| 1282 | num_embedding = solver->num_embedding; |
| 1283 | num_moment = solver->num_moment; |
| 1284 | num_negative = solver->num_negative; |
| 1285 | batch_size = solver->batch_size; |
| 1286 | optimizer = solver->optimizer; |
| 1287 | protocols = solver->protocols; |
| 1288 | sampler_protocol = solver->sampler_protocol; |
| 1289 | |
| 1290 | embeddings.resize(num_embedding); |
| 1291 | gradients.resize(num_embedding); |
| 1292 | moments.resize(num_embedding); |
| 1293 | for (int i = 0; i < num_embedding; i++) { |
| 1294 | Protocol protocol = protocols[i]; |
| 1295 | Index size = solver->embeddings[i]->size(); |
| 1296 | if (protocol & kHeadPartition) |
| 1297 | size = solver->head_partition_size; |
| 1298 | if (protocol & kTailPartition) |
| 1299 | size = solver->tail_partition_size; |
| 1300 | if (protocol & kSharedWithPredecessor && solver->num_partition == 1) |
| 1301 | size = 0; |
| 1302 | embeddings[i] = std::make_shared<Memory<Vector, Index>>(device_id, 0, work_stream); |
| 1303 | embeddings[i]->reallocate(size); |
| 1304 | gradients[i] = std::make_shared<Memory<Vector, Index>>(-1, 0, work_stream); |
| 1305 | if (!(protocol & kInPlace)) |
| 1306 | gradients[i]->reallocate(size); |
| 1307 | moments[i] = std::make_shared<std::vector<Memory<Vector, Index>>>(); |
| 1308 | for (int j = 0; j < num_moment; j++) { |
| 1309 | moments[i]->push_back(Memory<Vector, Index>(device_id, 0, work_stream)); |
| 1310 | (*moments[i])[j].reallocate(size); |
| 1311 | } |
| 1312 | } |
| 1313 | Index size = 0; |
| 1314 | if (sampler_protocol & kHeadPartition) |
| 1315 | size += solver->head_partition_size; |
| 1316 | if (sampler_protocol & kTailPartition) |
| 1317 | size += solver->tail_partition_size; |
| 1318 | if (sampler_protocol & kGlobal) |
| 1319 | size += solver->num_vertex; |
| 1320 | negative_sampler.reallocate(size); |
| 1321 | |
| 1322 | batch.resize(batch_size * kSampleSize); |
| 1323 | negative_batch.resize(batch_size * num_negative); |
| 1324 | loss.resize(batch_size); |
| 1325 | random.resize(batch_size * num_negative * 2); |
| 1326 | |
| 1327 | head_partition_id = -1; |
| 1328 | tail_partition_id = -1; |
| 1329 | } |
| 1330 | |
| 1331 | /** Free GPU memory */ |
| 1332 | void clear() { |
nothing calls this directly
no test coverage detected