From 74a83664dd9244bcd255e075235c5680e6174160 Mon Sep 17 00:00:00 2001 From: Jeff Donahue Date: Sat, 15 Feb 2014 16:28:44 -0800 Subject: [PATCH] allow in place computation of SplitLayer 0th top blob --- src/caffe/layers/split_layer.cpp | 28 +++++++++++++++-- src/caffe/test/test_split_layer.cpp | 48 +++++++++++++++++++++++++++++ 2 files changed, 74 insertions(+), 2 deletions(-) diff --git a/src/caffe/layers/split_layer.cpp b/src/caffe/layers/split_layer.cpp index cdfd53925c8..5accdd08e32 100644 --- a/src/caffe/layers/split_layer.cpp +++ b/src/caffe/layers/split_layer.cpp @@ -15,6 +15,12 @@ void SplitLayer::SetUp(const vector*>& bottom, CHECK_GE(top->size(), 1) << "Split Layer takes at least one blob as output."; count_ = bottom[0]->count(); for (int i = 0; i < top->size(); ++i) { + // Allow the 0th top blob to be 'in-place', but no others. + if (i == 0 && (*top)[i] == bottom[0]) { + continue; + } else { + CHECK_NE((*top)[i], bottom[0]) << "Only 0th top blob may be in place."; + } (*top)[i]->Reshape(bottom[0]->num(), bottom[0]->channels(), bottom[0]->height(), bottom[0]->width()); CHECK_EQ(count_, (*top)[i]->count()); @@ -26,6 +32,9 @@ void SplitLayer::Forward_cpu(const vector*>& bottom, vector*>* top) { const Dtype* bottom_data = bottom[0]->cpu_data(); for (int i = 0; i < top->size(); ++i) { + if (i == 0 && (*top)[i] == bottom[0]) { + continue; + } Dtype* top_data = (*top)[i]->mutable_cpu_data(); caffe_copy(count_, bottom_data, top_data); } @@ -36,6 +45,9 @@ void SplitLayer::Forward_gpu(const vector*>& bottom, vector*>* top) { const Dtype* bottom_data = bottom[0]->gpu_data(); for (int i = 0; i < top->size(); ++i) { + if (i == 0 && (*top)[i] == bottom[0]) { + continue; + } Dtype* top_data = (*top)[i]->mutable_gpu_data(); caffe_gpu_copy(count_, bottom_data, top_data); } @@ -47,7 +59,13 @@ Dtype SplitLayer::Backward_cpu(const vector*>& top, if (propagate_down) { const Dtype* top_diff = top[0]->cpu_diff(); Dtype* bottom_diff = (*bottom)[0]->mutable_cpu_diff(); - caffe_copy(count_, top_diff, bottom_diff); + // Initialize by copying first top blob diff to our diff, unless we're + // doing in-place computation for the first blob, in which case the diff is + // already initialized. + if (top[0] != (*bottom)[0]) { + caffe_copy(count_, top_diff, bottom_diff); + } + // Add remaining top blob diffs. for (int i = 1; i < top.size(); ++i) { top_diff = top[i]->cpu_diff(); caffe_axpy(count_, Dtype(1.), top_diff, bottom_diff); @@ -63,7 +81,13 @@ Dtype SplitLayer::Backward_gpu(const vector*>& top, if (propagate_down) { const Dtype* top_diff = top[0]->gpu_diff(); Dtype* bottom_diff = (*bottom)[0]->mutable_gpu_diff(); - caffe_gpu_copy(count_, top_diff, bottom_diff); + // Initialize by copying first top blob diff to our diff, unless we're + // doing in-place computation for the first blob, in which case the diff is + // already initialized. + if (top[0] != (*bottom)[0]) { + caffe_gpu_copy(count_, top_diff, bottom_diff); + } + // Add remaining top blob diffs. for (int i = 1; i < top.size(); ++i) { top_diff = top[i]->gpu_diff(); caffe_gpu_axpy(count_, Dtype(1.), top_diff, bottom_diff); diff --git a/src/caffe/test/test_split_layer.cpp b/src/caffe/test/test_split_layer.cpp index a3252b23378..f6d5d521fec 100644 --- a/src/caffe/test/test_split_layer.cpp +++ b/src/caffe/test/test_split_layer.cpp @@ -88,6 +88,32 @@ TYPED_TEST(SplitLayerTest, TestGPU) { } } +TYPED_TEST(SplitLayerTest, TestCPUInPlace) { + LayerParameter layer_param; + SplitLayer layer(layer_param); + Caffe::set_mode(Caffe::CPU); + this->blob_top_vec_[0] = this->blob_bottom_vec_[0]; + layer.SetUp(this->blob_bottom_vec_, &(this->blob_top_vec_)); + layer.Forward(this->blob_bottom_vec_, &(this->blob_top_vec_)); + for (int i = 0; i < this->blob_bottom_->count(); ++i) { + TypeParam bottom_value = this->blob_bottom_->cpu_data()[i]; + EXPECT_EQ(bottom_value, this->blob_top_b_->cpu_data()[i]); + } +} + +TYPED_TEST(SplitLayerTest, TestGPUInPlace) { + LayerParameter layer_param; + SplitLayer layer(layer_param); + Caffe::set_mode(Caffe::GPU); + this->blob_top_vec_[0] = this->blob_bottom_vec_[0]; + layer.SetUp(this->blob_bottom_vec_, &(this->blob_top_vec_)); + layer.Forward(this->blob_bottom_vec_, &(this->blob_top_vec_)); + for (int i = 0; i < this->blob_bottom_->count(); ++i) { + TypeParam bottom_value = this->blob_bottom_->cpu_data()[i]; + EXPECT_EQ(bottom_value, this->blob_top_b_->cpu_data()[i]); + } +} + TYPED_TEST(SplitLayerTest, TestCPUGradient) { LayerParameter layer_param; Caffe::set_mode(Caffe::CPU); @@ -108,6 +134,28 @@ TYPED_TEST(SplitLayerTest, TestGPUGradient) { this->blob_top_vec_); } +TYPED_TEST(SplitLayerTest, TestCPUGradientInPlace) { + LayerParameter layer_param; + Caffe::set_mode(Caffe::CPU); + SplitLayer layer(layer_param); + GradientChecker checker(1e-2, 1e-2); + this->blob_top_vec_[0] = this->blob_bottom_vec_[0]; + checker.CheckGradientExhaustive(layer, this->blob_bottom_vec_, this->blob_top_vec_); + checker.CheckGradientExhaustive(layer, this->blob_bottom_vec_, + this->blob_top_vec_); +} + +TYPED_TEST(SplitLayerTest, TestGPUGradientInPlace) { + LayerParameter layer_param; + Caffe::set_mode(Caffe::GPU); + SplitLayer layer(layer_param); + GradientChecker checker(1e-2, 1e-2); + this->blob_top_vec_[0] = this->blob_bottom_vec_[0]; + checker.CheckGradientExhaustive(layer, this->blob_bottom_vec_, this->blob_top_vec_); + checker.CheckGradientExhaustive(layer, this->blob_bottom_vec_, + this->blob_top_vec_); +} + template class SplitLayerInsertionTest : public ::testing::Test {