* @brief Train embeddings with samples in the sample block * @param _head_partition_id id of head partition * @param _tail_partition_id id of tail partition */
| 1509 | * @param _tail_partition_id id of tail partition |
| 1510 | */ |
| 1511 | void train(int _head_partition_id, int _tail_partition_id) { |
| 1512 | CUDA_CHECK(cudaSetDevice(device_id)); |
| 1513 | load_partition(_head_partition_id, _tail_partition_id); |
| 1514 | |
| 1515 | auto &samples = solver->sample_pools[solver->pool_id][head_partition_id][tail_partition_id]; |
| 1516 | log_frequency = solver->log_frequency; |
| 1517 | for (int i = 0; i < solver->positive_reuse; i++) |
| 1518 | for (int j = 0; j < solver->episode_size; j++) { |
| 1519 | batch.copy(&samples[j * batch_size], batch_size * kSampleSize); |
| 1520 | train_batch(solver->batch_id++); |
| 1521 | } |
| 1522 | } |
| 1523 | |
| 1524 | /** Train a single batch */ |
| 1525 | virtual void train_batch(int batch_id) { |