| 161 | } |
| 162 | |
| 163 | void ggml_opt_dataset_get_batch(ggml_opt_dataset_t dataset, struct ggml_tensor * data_batch, struct ggml_tensor * labels_batch, int64_t ibatch) { |
| 164 | GGML_ASSERT( data_batch && ggml_is_contiguous(data_batch)); |
| 165 | GGML_ASSERT(!labels_batch || ggml_is_contiguous(labels_batch)); |
| 166 | GGML_ASSERT((labels_batch == nullptr) == (dataset->labels == nullptr)); |
| 167 | GGML_ASSERT( data_batch->type == dataset->data->type); |
| 168 | GGML_ASSERT(!labels_batch || labels_batch->type == dataset->labels->type); |
| 169 | |
| 170 | const size_t nb_data_batch = ggml_nbytes(data_batch); |
| 171 | GGML_ASSERT(nb_data_batch % dataset->nbs_data == 0); |
| 172 | const int64_t shards_per_batch = nb_data_batch / dataset->nbs_data; |
| 173 | |
| 174 | if (labels_batch) { |
| 175 | const size_t nb_labels_batch = ggml_nbytes(labels_batch); |
| 176 | GGML_ASSERT(nb_labels_batch == shards_per_batch*dataset->nbs_labels); |
| 177 | } |
| 178 | |
| 179 | GGML_ASSERT((ibatch + 1)*shards_per_batch <= int64_t(dataset->permutation.size())); |
| 180 | |
| 181 | for (int64_t ishard_batch = 0; ishard_batch < shards_per_batch; ++ishard_batch) { |
| 182 | const int64_t ishard = dataset->permutation[ibatch*shards_per_batch + ishard_batch]; |
| 183 | |
| 184 | const char * ptr_data = (const char *) dataset->data->data + ishard*dataset->nbs_data; |
| 185 | ggml_backend_tensor_set(data_batch, ptr_data, ishard_batch*dataset->nbs_data, dataset->nbs_data); |
| 186 | |
| 187 | if (!labels_batch) { |
| 188 | continue; |
| 189 | } |
| 190 | |
| 191 | const char * ptr_labels = (const char *) dataset->labels->data + ishard*dataset->nbs_labels; |
| 192 | ggml_backend_tensor_set(labels_batch, ptr_labels, ishard_batch*dataset->nbs_labels, dataset->nbs_labels); |
| 193 | } |
| 194 | } |
| 195 | |
| 196 | void ggml_opt_dataset_get_batch_host(ggml_opt_dataset_t dataset, void * data_batch, size_t nb_data_batch, void * labels_batch, int64_t ibatch) { |
| 197 | GGML_ASSERT((labels_batch == nullptr) == (dataset->labels == nullptr)); |