mirror of https://github.com/davisking/dlib.git
Semantic-segmentation loss calculation: fix buffer usage on multi-GPU training (#1717)
* Semantic-segmentation loss calculation: fix buffer usage on multi-GPU training * Review fix: make the work buffer live longer
This commit is contained in:
parent
9433bfd6ab
commit
ccd8b64f1b
|
@ -423,7 +423,6 @@ namespace dlib
|
|||
compute_loss_multiclass_log_per_pixel(
|
||||
)
|
||||
{
|
||||
work = device_global_buffer();
|
||||
}
|
||||
|
||||
template <
|
||||
|
@ -439,6 +438,10 @@ namespace dlib
|
|||
const size_t bytes_per_plane = subnetwork_output.nr()*subnetwork_output.nc()*sizeof(uint16_t);
|
||||
// Allocate a cuda buffer to store all the truth images and also one float
|
||||
// for the scalar loss output.
|
||||
if (!work)
|
||||
{
|
||||
work = device_global_buffer();
|
||||
}
|
||||
cuda_data_void_ptr buf = work->get(subnetwork_output.num_samples()*bytes_per_plane + sizeof(float));
|
||||
|
||||
cuda_data_void_ptr loss_buf = buf;
|
||||
|
@ -467,7 +470,7 @@ namespace dlib
|
|||
double& loss
|
||||
);
|
||||
|
||||
std::shared_ptr<resizable_cuda_buffer> work;
|
||||
mutable std::shared_ptr<resizable_cuda_buffer> work;
|
||||
};
|
||||
|
||||
// ------------------------------------------------------------------------------------
|
||||
|
|
Loading…
Reference in New Issue