Added support for all the other dlib layers that make sense.

tools/convert_dlib_nets_to_caffe/main.cpp

@@ -14,7 +14,7 @@ using namespace dlib;
 // ----------------------------------------------------------------------------------------
 
 // Only these computational layers have parameters
-const std::set<string> comp_tags_with_params = {"fc", "fc_no_bias", "con", "bn_con", "bn_fc", "affine", "prelu"};
+const std::set<string> comp_tags_with_params = {"fc", "fc_no_bias", "con", "affine_con", "affine_fc", "affine", "prelu"};
 
 struct layer
 {
@@ -28,6 +28,15 @@ struct layer
     long skip_id = -1;  // If this isn't -1 then it means this layer draws its inputs from
                         // the most recent layer with tag_id==skip_id rather than its immediate predecessor. 
 
+    double attribute (const string& key) const
+    {
+        auto i = attributes.find(key);
+        if (i != attributes.end())
+            return i->second;
+        else
+            throw dlib::error("Layer doesn't have the requested attribute '" + key + "'.");
+    }
+
     string caffe_layer_name() const 
     { 
         if (type == "input")
@@ -71,10 +80,10 @@ string find_layer_caffe_name (
         {
             i--;
             // if we hit the end of the network before we found what we were looking for
-            if (i->type == "input")
-                throw dlib::error("Network definition is bad, a layer wanted to skip back to a non-existing layer.");
             if (i->tag_id == tag_id)
                 return i->caffe_layer_name();
+            if (i->type == "input")
+                throw dlib::error("Network definition is bad, a layer wanted to skip back to a non-existing layer.");
         }
     }
 }
@@ -99,30 +108,30 @@ void convert_dlib_xml_to_cafffe_python_code(
     const string& xml_filename
 )
 {
-    auto layers = parse_dlib_xml(xml_filename);
+    const auto layers = parse_dlib_xml(xml_filename);
 
     cout << "import caffe " << endl;
     cout << "from caffe import layers as L, params as P" << endl;
     cout << "import numpy as np" << endl;
 
-    // dlib nets don't commit to a batch size, so just use 32 as the default
-    cout << "batch_size = 32;" << endl;
+    // dlib nets don't commit to a batch size, so just use 1 as the default
+    cout << "batch_size = 1;" << endl;
     if (layers.back().detail_name == "input_rgb_image")
     {
-        cout << "input_nr = 150; #WARNING, the source dlib network didn't commit to a specific input size, so we put 150 here as a default." << endl;
-        cout << "input_nc = 150; #WARNING, the source dlib network didn't commit to a specific input size, so we put 150 here as a default." << endl;
+        cout << "input_nr = 28; #WARNING, the source dlib network didn't commit to a specific input size, so we put 28 here as a default." << endl;
+        cout << "input_nc = 28; #WARNING, the source dlib network didn't commit to a specific input size, so we put 28 here as a default." << endl;
         cout << "input_k = 3;" << endl;
     }
     else if (layers.back().detail_name == "input_rgb_image_sized")
     {
-        cout << "input_nr = " << layers.back().attributes["nr"] << ";" << endl;
-        cout << "input_nc = " << layers.back().attributes["nc"] << ";" << endl;
+        cout << "input_nr = " << layers.back().attribute("nr") << ";" << endl;
+        cout << "input_nc = " << layers.back().attribute("nc") << ";" << endl;
         cout << "input_k = 3;" << endl;
     }
     else if (layers.back().detail_name == "input")
     {
-        cout << "input_nr = 150; #WARNING, the source dlib network didn't commit to a specific input size, so we put 150 here as a default." << endl;
-        cout << "input_nc = 150; #WARNING, the source dlib network didn't commit to a specific input size, so we put 150 here as a default." << endl;
+        cout << "input_nr = 28; #WARNING, the source dlib network didn't commit to a specific input size, so we put 28 here as a default." << endl;
+        cout << "input_nc = 28; #WARNING, the source dlib network didn't commit to a specific input size, so we put 28 here as a default." << endl;
         cout << "input_k = 1;" << endl;
     }
     else
@@ -131,6 +140,8 @@ void convert_dlib_xml_to_cafffe_python_code(
     }
 
     cout << "def make_netspec():" << endl;
+    cout << "    # For reference, the only \"documentation\" about caffe layer parameters seems to be this page:\n";
+    cout << "    # https://github.com/BVLC/caffe/blob/master/src/caffe/proto/caffe.proto\n" << endl;
     cout << "    n = caffe.NetSpec(); " << endl;
     cout << "    n.data,n.label = L.MemoryData(batch_size=batch_size, channels=input_k, height=input_nr, width=input_nc, ntop=2)" << endl;
     // iterate the layers starting with the input layer
@@ -144,13 +155,13 @@ void convert_dlib_xml_to_cafffe_python_code(
         if (i->detail_name == "con")
         {
             cout << "    n." << i->caffe_layer_name() << " = L.Convolution(n." << find_input_layer_caffe_name(i);
-            cout << ", num_output=" << i->attributes["num_filters"];
-            cout << ", kernel_w=" << i->attributes["nc"];
-            cout << ", kernel_h=" << i->attributes["nr"];
-            cout << ", stride_w=" << i->attributes["stride_x"];
-            cout << ", stride_h=" << i->attributes["stride_y"];
-            cout << ", pad_w=" << i->attributes["padding_x"];
-            cout << ", pad_h=" << i->attributes["padding_y"];
+            cout << ", num_output=" << i->attribute("num_filters");
+            cout << ", kernel_w=" << i->attribute("nc");
+            cout << ", kernel_h=" << i->attribute("nr");
+            cout << ", stride_w=" << i->attribute("stride_x");
+            cout << ", stride_h=" << i->attribute("stride_y");
+            cout << ", pad_w=" << i->attribute("padding_x");
+            cout << ", pad_h=" << i->attribute("padding_y");
             cout << ");\n";
         }
         else if (i->detail_name == "relu")
@@ -162,51 +173,91 @@ void convert_dlib_xml_to_cafffe_python_code(
         {
             cout << "    n." << i->caffe_layer_name() << " = L.Pooling(n." << find_input_layer_caffe_name(i);
             cout << ", pool=P.Pooling.MAX"; 
-            cout << ", kernel_w=" << i->attributes["nc"];
-            cout << ", kernel_h=" << i->attributes["nr"];
-            cout << ", stride_w=" << i->attributes["stride_x"];
-            cout << ", stride_h=" << i->attributes["stride_y"];
-            cout << ", pad_w=" << i->attributes["padding_x"];
-            cout << ", pad_h=" << i->attributes["padding_y"];
+            if (i->attribute("nc")==0)
+            {
+                cout << ", global_pooling=True";
+            }
+            else
+            {
+                cout << ", kernel_w=" << i->attribute("nc");
+                cout << ", kernel_h=" << i->attribute("nr");
+            }
+            if (i->attribute("padding_x") != 0 || i->attribute("padding_y") != 0)
+            {
+                throw dlib::error("dlib and caffe implement pooling with non-zero padding differently, so you can't convert a "
+                    "network with such pooling layers.");
+            }
+
+            cout << ", stride_w=" << i->attribute("stride_x");
+            cout << ", stride_h=" << i->attribute("stride_y");
+            cout << ", pad_w=" << i->attribute("padding_x");
+            cout << ", pad_h=" << i->attribute("padding_y");
             cout << ");\n";
         }
         else if (i->detail_name == "avg_pool")
         {
             cout << "    n." << i->caffe_layer_name() << " = L.Pooling(n." << find_input_layer_caffe_name(i);
-            cout << ", pool=P.Pooling.MAX"; 
-            cout << ", kernel_w=" << i->attributes["nc"];
-            cout << ", kernel_h=" << i->attributes["nr"];
-            cout << ", stride_w=" << i->attributes["stride_x"];
-            cout << ", stride_h=" << i->attributes["stride_y"];
-            cout << ", pad_w=" << i->attributes["padding_x"];
-            cout << ", pad_h=" << i->attributes["padding_y"];
+            cout << ", pool=P.Pooling.AVE"; 
+            if (i->attribute("nc")==0)
+            {
+                cout << ", global_pooling=True";
+            }
+            else
+            {
+                cout << ", kernel_w=" << i->attribute("nc");
+                cout << ", kernel_h=" << i->attribute("nr");
+            }
+            if (i->attribute("padding_x") != 0 || i->attribute("padding_y") != 0)
+            {
+                throw dlib::error("dlib and caffe implement pooling with non-zero padding differently, so you can't convert a "
+                    "network with such pooling layers.");
+            }
+
+            cout << ", stride_w=" << i->attribute("stride_x");
+            cout << ", stride_h=" << i->attribute("stride_y");
+            cout << ", pad_w=" << i->attribute("padding_x");
+            cout << ", pad_h=" << i->attribute("padding_y");
             cout << ");\n";
         }
         else if (i->detail_name == "fc")
         {
             cout << "    n." << i->caffe_layer_name() << " = L.InnerProduct(n." << find_input_layer_caffe_name(i);
-            cout << ", num_output=" << i->attributes["num_outputs"];
+            cout << ", num_output=" << i->attribute("num_outputs");
             cout << ", bias_term=True";
             cout << ");\n";
         }
         else if (i->detail_name == "fc_no_bias")
         {
             cout << "    n." << i->caffe_layer_name() << " = L.InnerProduct(n." << find_input_layer_caffe_name(i);
-            cout << ", num_output=" << i->attributes["num_outputs"];
+            cout << ", num_output=" << i->attribute("num_outputs");
             cout << ", bias_term=False";
             cout << ");\n";
         }
-        else if (i->detail_name == "bn_con")
+        else if (i->detail_name == "bn_con" || i->detail_name == "bn_fc")
         {
-            // TODO
+            throw dlib::error("Conversion from dlib's batch norm layers to caffe's isn't supported.  Instead, "
+                "you should put your network into 'test mode' by switching batch norm layers to affine layers.");
         }
-        else if (i->detail_name == "bn_fc")
+        else if (i->detail_name == "affine_con")
         {
-            // TODO
+            cout << "    n." << i->caffe_layer_name() << " = L.Scale(n." << find_input_layer_caffe_name(i);
+            cout << ", axis=1";
+            cout << ", bias_term=True";
+            cout << ");\n";
+        }
+        else if (i->detail_name == "affine_fc")
+        {
+            cout << "    n." << i->caffe_layer_name() << " = L.Scale(n." << find_input_layer_caffe_name(i);
+            cout << ", axis=3";
+            cout << ", bias_term=True";
+            cout << ");\n";
         }
         else if (i->detail_name == "add_prev")
         {
-            // TODO
+            cout << "    n." << i->caffe_layer_name() << " = L.Eltwise(n." << find_input_layer_caffe_name(i);
+            cout << ", n." << find_layer_caffe_name(i, i->attribute("tag"));
+            cout << ", operation=P.Eltwise.SUM";
+            cout << ");\n";
         }
         else
         {
@@ -215,6 +266,11 @@ void convert_dlib_xml_to_cafffe_python_code(
     }
     cout << "    return n.to_proto();\n\n" << endl;
 
+
+    // -------------------------
+    // -------------------------
+
+
     cout << "def save_as_caffe_model(def_file, weights_file):\n";
     cout << "    with open(def_file, 'w') as f: f.write(str(make_netspec()));\n";
     cout << "    net = caffe.Net(def_file, caffe.TEST);\n";
@@ -222,6 +278,9 @@ void convert_dlib_xml_to_cafffe_python_code(
     cout << "    net.save(weights_file);\n\n";
 
 
+    // -------------------------
+    // -------------------------
+
 
     cout << "def set_network_weights(net):\n";
     cout << "    # populate network parameters\n";
@@ -235,7 +294,7 @@ void convert_dlib_xml_to_cafffe_python_code(
 
         if (i->detail_name == "con")
         {
-            const long num_filters = i->attributes["num_filters"];
+            const long num_filters = i->attribute("num_filters");
             matrix<double> weights = trans(rowm(i->params,range(0,i->params.size()-num_filters-1)));
             matrix<double> biases  = trans(rowm(i->params,range(i->params.size()-num_filters, i->params.size()-1)));
 
@@ -273,13 +332,21 @@ void convert_dlib_xml_to_cafffe_python_code(
             cout << "    p.shape = net.params['"<<i->caffe_layer_name()<<"'][0].data.shape;\n";
             cout << "    net.params['"<<i->caffe_layer_name()<<"'][0].data[:] = p;\n";
         }
-        else if (i->detail_name == "bn_con")
+        else if (i->detail_name == "affine_con" || i->detail_name == "affine_fc")
         {
-            // TODO
-        }
-        else if (i->detail_name == "bn_fc")
-        {
-            // TODO
+            const long dims = i->params.size()/2;
+            matrix<double> gamma = trans(rowm(i->params,range(0,dims-1)));
+            matrix<double> beta  = trans(rowm(i->params,range(dims, 2*dims-1)));
+
+            // set gamma weights
+            cout << "    p = "; print_as_np_array(cout,gamma); cout << ";\n";
+            cout << "    p.shape = net.params['"<<i->caffe_layer_name()<<"'][0].data.shape;\n";
+            cout << "    net.params['"<<i->caffe_layer_name()<<"'][0].data[:] = p;\n";
+
+            // set beta weights 
+            cout << "    p = "; print_as_np_array(cout,beta); cout << ";\n";
+            cout << "    p.shape = net.params['"<<i->caffe_layer_name()<<"'][1].data.shape;\n";
+            cout << "    net.params['"<<i->caffe_layer_name()<<"'][1].data[:] = p;\n";
         }
     }
 
@@ -289,6 +356,7 @@ void convert_dlib_xml_to_cafffe_python_code(
 
 int main(int argc, char** argv) try
 {
+    cout.precision(9);
     // TODO, write out to multiple files or just process one file at a time.  
     for (int i = 1; i < argc; ++i)
         convert_dlib_xml_to_cafffe_python_code(argv[i]);