try to fix it reusing the code... not sure though

This commit is contained in:
Adrià Arrufat 2021-09-26 20:47:00 +09:00
parent 4851d54f0f
commit f308cac6df
1 changed files with 7 additions and 39 deletions

View File

@ -1194,6 +1194,9 @@ namespace dlib
job.have_data.resize(devs); job.have_data.resize(devs);
job.test_only = test_only; job.test_only = test_only;
// check if the iterator points to anything
const bool have_labels = sizeof(*lbegin) != 1;
// chop the data into devs blocks, each of about block_size elements. // chop the data into devs blocks, each of about block_size elements.
const double block_size = num / static_cast<double>(devs); const double block_size = num / static_cast<double>(devs);
@ -1211,7 +1214,8 @@ namespace dlib
if (start < stop) if (start < stop)
{ {
devices[i]->net.to_tensor(dbegin+start, dbegin+stop, job.t[i]); devices[i]->net.to_tensor(dbegin+start, dbegin+stop, job.t[i]);
job.labels[i].assign(lbegin+start, lbegin+stop); if (have_labels)
job.labels[i].assign(lbegin+start, lbegin+stop);
job.have_data[i] = true; job.have_data[i] = true;
} }
else else
@ -1237,44 +1241,8 @@ namespace dlib
data_iterator dend data_iterator dend
) )
{ {
propagate_exception(); typename std::vector<training_label_type>::iterator nothing;
size_t num = std::distance(dbegin, dend); send_job(test_only, dbegin, dend, nothing);
size_t devs = devices.size();
job.t.resize(devs);
job.have_data.resize(devs);
job.test_only = test_only;
// chop the data into devs blocks, each of about block_size elements.
const double block_size = num / static_cast<double>(devs);
const auto prev_dev = dlib::cuda::get_device();
double j = 0;
for (size_t i = 0; i < devs; ++i)
{
dlib::cuda::set_device(devices[i]->device_id);
const size_t start = static_cast<size_t>(std::round(j));
const size_t stop = static_cast<size_t>(std::round(j + block_size));
if (start < stop)
{
devices[i]->net.to_tensor(dbegin+start, dbegin+stop, job.t[i]);
job.have_data[i] = true;
}
else
{
job.have_data[i] = false;
}
j += block_size;
}
DLIB_ASSERT(std::fabs(j - num) < 1e-10);
dlib::cuda::set_device(prev_dev);
job_pipe.enqueue(job);
} }
void print_progress() void print_progress()