From 6f36b7f554f0fff62dd7688d456cd4d0274fc3a8 Mon Sep 17 00:00:00 2001 From: Yangqing Jia Date: Thu, 24 Oct 2013 10:47:46 -0700 Subject: [PATCH] data layer race condition bugfix --- examples/demo_mnist.cpp | 1 + src/caffe/layers/data_layer.cpp | 5 +++-- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/examples/demo_mnist.cpp b/examples/demo_mnist.cpp index 9eb6204..11d3fc5 100644 --- a/examples/demo_mnist.cpp +++ b/examples/demo_mnist.cpp @@ -25,6 +25,7 @@ int main(int argc, char** argv) { return 0; } google::InitGoogleLogging(argv[0]); + Caffe::DeviceQuery(); if (argc == 4 && strcmp(argv[3], "GPU") == 0) { LOG(ERROR) << "Using GPU"; diff --git a/src/caffe/layers/data_layer.cpp b/src/caffe/layers/data_layer.cpp index fd26c5d..9ed9516 100644 --- a/src/caffe/layers/data_layer.cpp +++ b/src/caffe/layers/data_layer.cpp @@ -170,8 +170,9 @@ void DataLayer::SetUp(const vector*>& bottom, // cpu_data calls so that the prefetch thread does not accidentally make // simultaneous cudaMalloc calls when the main thread is running. In some // GPUs this seems to cause failures if we do not so. - layer->prefetch_data_->mutable_cpu_data(); - layer->prefetch_label_->mutable_cpu_data(); + prefetch_data_->mutable_cpu_data(); + prefetch_label_->mutable_cpu_data(); + data_mean_.cpu_data(); // LOG(INFO) << "Initializing prefetch"; CHECK(!pthread_create(&thread_, NULL, DataLayerPrefetch, reinterpret_cast(this))) << "Pthread execution failed."; -- 2.7.4