void MemCopy::compute_cpu(const vector<bool>& add) { std::lock_guard<std::mutex>lock_guard_(outputs_[0]->get_mutex()); if (inputs_[0]->cpu_data() == outputs_[0]->mutable_cpu_data()) { return; } else { caffe::caffe_cpu_copy<DTYPE>(inputs_[0]->size().count(), inputs_[0]->cpu_data(), outputs_[0]->mutable_cpu_data()); } }
//返回异常, 主要因为容器中的元素可能为指针, 也可能为基本类型 如 int, 则在为空的情况下返回 NULL 或 0都不合适, 所以返回值改为用户指定值,可以由用户指定为0或NULL T take(T empty_return_value) { boost::lock_guard<boost::mutex> lock_guard_(_mutex); if (_container.size() == 0) { return empty_return_value; } T front = _container.front(); _container.pop(); return front; }
void MemCopy::compute_gpu(const vector<bool>& add) { std::lock_guard<std::mutex>lock_guard_(outputs_[0]->get_mutex()); if (inputs_[0]->device() >= 0 && outputs_[0]->device() >=0 && inputs_[0]->device() == outputs_[0]->device() && inputs_[0]->gpu_data() == outputs_[0]->mutable_gpu_data()) { return; } const DTYPE* src = inputs_[0]->device() < 0 ? inputs_[0]->cpu_data() : inputs_[0]->gpu_data();; DTYPE* dst = outputs_[0]->device() < 0 ? outputs_[0]->mutable_cpu_data() : outputs_[0]->mutable_gpu_data(); CUDA_CHECK(cudaMemcpyAsync(dst, src, inputs_[0]->size().count() * sizeof(DTYPE), cudaMemcpyDefault, stream())); }
size_t size() { boost::lock_guard<boost::mutex> lock_guard_(_mutex); return _container.size(); }
void put(const T& task) { boost::lock_guard<boost::mutex> lock_guard_(_mutex); _container.push(task); }