Make device factory usage more inline with layer factory.

Author: Rob Hess

include/caffe/common.hpp

@@ -105,7 +105,7 @@ class Caffe {
     }
     return *singleton_;
   }
-  enum Brew { CPU, GPU };
+  enum Brew { UNSPECIFIED = -1, CPU, GPU };
   enum Phase { TRAIN, TEST };
 
 

include/caffe/device.hpp

@@ -8,6 +8,7 @@
 
 #include "glog/logging.h"
 
+#include "caffe/common.hpp"
 #include "caffe/util/im2col.hpp"
 #include "caffe/util/math_functions.hpp"
 
@@ -252,14 +253,9 @@ class GPUDevice : public Device<Dtype> {
       const int stride, Dtype* data_im);
 };
 
+// Device factory function
 template<typename Dtype>
-class DeviceFactory {
- public:
-  static Device<Dtype>* GetDevice();
- private:
-  static Device<Dtype>* cpu_device_;
-  static Device<Dtype>* gpu_device_;
-};
+Device<Dtype>* GetDevice(Caffe::Brew mode = Caffe::UNSPECIFIED);
 
 }  // namespace caffe
 

include/caffe/layer.hpp

@@ -19,7 +19,7 @@ class Layer {
   // to SetUp(), where the dimensions of the bottom blobs are provided to the
   // layer.
   explicit Layer(const LayerParameter& param)
-    : layer_param_(param), device_(DeviceFactory<Dtype>::GetDevice()) {
+    : layer_param_(param), device_(GetDevice<Dtype>()) {
       // The only thing we do is to copy blobs if there are any.
       if (layer_param_.blobs_size() > 0) {
         blobs_.resize(layer_param_.blobs_size());

src/caffe/device_factory.cpp

@@ -0,0 +1,51 @@
+// Copyright 2014 BVLC and contributors.
+
+#include "caffe/common.hpp"
+#include "caffe/device.hpp"
+
+namespace caffe {
+
+template<typename Dtype>
+class DeviceFactory {
+ public:
+  static Device<Dtype>* GetDevice(Caffe::Brew mode);
+ private:
+  static Device<Dtype>* cpu_device_;
+  static Device<Dtype>* gpu_device_;
+};
+
+template<typename Dtype>
+Device<Dtype>* DeviceFactory<Dtype>::GetDevice(Caffe::Brew mode) {
+  if (mode == Caffe::UNSPECIFIED)
+    mode = Caffe::mode();
+  switch (mode) {
+    case Caffe::CPU:
+      return cpu_device_;
+    case Caffe::GPU:
+      return gpu_device_;
+    default:
+      LOG(FATAL) << "Unknown caffe mode.";
+      return static_cast<Device<Dtype>*>(NULL);
+  }
+}
+
+template<typename Dtype>
+Device<Dtype>* DeviceFactory<Dtype>::cpu_device_ = new CPUDevice<Dtype>();
+
+template<typename Dtype>
+Device<Dtype>* DeviceFactory<Dtype>::gpu_device_ = new GPUDevice<Dtype>();
+
+INSTANTIATE_CLASS(DeviceFactory);
+
+// A function to get the device either for the current mode (if
+// Caffe::UNSPECIFIED is passed, which is the default), or for a specific mode
+// (if a specific device's Caffe::Brew is passed).
+template<typename Dtype>
+Device<Dtype>* GetDevice(Caffe::Brew mode) {
+  return DeviceFactory<Dtype>::GetDevice(mode);
+}
+
+template Device<float>* GetDevice(Caffe::Brew mode);
+template Device<double>* GetDevice(Caffe::Brew mode);
+
+}  // namespace caffe

src/caffe/net.cpp

@@ -684,7 +684,7 @@ void Net<Dtype>::Update() {
     const int count = params_[i]->count();
     const Dtype* this_diff = params_[i]->const_diff();
     Dtype* owner_diff = params_[param_owners_[i]]->mutable_diff();
-    Device<Dtype>* device = DeviceFactory<Dtype>::GetDevice();
+    Device<Dtype>* device = GetDevice<Dtype>();
     device->add(count, this_diff, owner_diff, owner_diff);
   }
   // Now, update the owned parameters.

src/caffe/solver.cpp

@@ -367,7 +367,7 @@ void SGDSolver<Dtype>::ComputeUpdateValue() {
   }
   Dtype momentum = this->param_.momentum();
   Dtype weight_decay = this->param_.weight_decay();
-  Device<Dtype>* device = DeviceFactory<Dtype>::GetDevice();
+  Device<Dtype>* device = GetDevice<Dtype>();
   for (int param_id = 0; param_id < net_params.size(); ++param_id) {
     // Compute the value to history, and then copy them to the blob's diff.
     Dtype local_rate = rate * net_params_lr[param_id];