CaffeTransform is TimeVarying on GPU and not otherwise

Jordan Cheney
1 parent dff3e093
Showing 1 changed file with 46 additions and 10 deletions
openbr/plugins/classification/caffe.cpp
@@ -16,11 +16,46 @@ using namespace cv;
 namespace br
 {
+// Net doesn't expose a default constructor which is expected by the default resource allocator.
+// To get around that we make this custom stub class which has a default constructor that passes
+// empty values to the Net constructor.
+class CaffeNet : public Net<float>
+{
+public:
+    CaffeNet() : Net<float>("", caffe::TEST) {}
+    CaffeNet(const QString &model, caffe::Phase phase) : Net<float>(model.toStdString(), phase) {}
+};
+
+class CaffeResourceMaker : public ResourceMaker<CaffeNet>
+{
+    QString model;
+    QString weights;
+    int gpuDevice;
+
+public:
+    CaffeResourceMaker(const QString &model, const QString &weights, int gpuDevice) : model(model), weights(weights), gpuDevice(gpuDevice) {}
+
+private:
+    CaffeNet *make() const
+    {
+        if (gpuDevice >= 0) {
+            Caffe::SetDevice(gpuDevice);
+            Caffe::set_mode(Caffe::GPU);
+        } else {
+            Caffe::set_mode(Caffe::CPU);
+        }
+
+        CaffeNet *net = new CaffeNet(model, caffe::TEST);
+        net->CopyTrainedLayersFrom(weights.toStdString());
+        return net;
+    }
+};
+
 /*!
  * \brief A transform that wraps the Caffe deep learning library. This transform expects the input to a given Caffe model to be a MemoryDataLayer.
  * The output of the Caffe network is treated as a feature vector and is stored in dst. Batch processing is possible. For a given batch size set in
  * the memory data layer, src is expected to have an equal number of mats. Dst will always have the same size (number of mats) as src and the ordering
- * will be preserved, so dst[1] is the output of src[1] after it passes through the
+ * will be preserved, so dst[1] is the output of src[1] after it passes through the neural net.
  * \author Jordan Cheney \cite jcheney
  * \br_property QString model path to prototxt model file
  * \br_property QString weights path to caffemodel file
@@ -39,23 +74,22 @@ class CaffeFVTransform : public UntrainableTransform
     BR_PROPERTY(QString, weights, "")
     BR_PROPERTY(int, gpuDevice, -1)
-    QSharedPointer<Net<float> > net;
+    Resource<CaffeNet> caffeResource;
     void init()
     {
-        if (gpuDevice >= 0) {
-            Caffe::SetDevice(gpuDevice);
-            Caffe::set_mode(Caffe::GPU);
-        } else {
-            Caffe::set_mode(Caffe::CPU);
-        }
+        caffeResource.setResourceMaker(new CaffeResourceMaker(model, weights, gpuDevice));
+    }
-        net.reset(new Net<float>(model.toStdString(), caffe::TEST));
-        net->CopyTrainedLayersFrom(weights.toStdString());
+    bool timeVarying() const
+    {
+        return gpuDevice < 0 ? false : true;
     }
     void project(const Template &src, Template &dst) const
     {
+        CaffeNet *net = caffeResource.acquire();
+
         MemoryDataLayer<float> *data_layer = static_cast<MemoryDataLayer<float> *>(net->layers()[0].get());
         if (src.size() != data_layer->batch_size())
@@ -70,6 +104,8 @@ class CaffeFVTransform : public UntrainableTransform
         int dim_features = output->count() / data_layer->batch_size();
         for (int n = 0; n < data_layer->batch_size(); n++)
             dst += Mat(1, dim_features, CV_32FC1, output->mutable_cpu_data() + output->offset(n));
+
+        caffeResource.release(net);
     }
 };