Example #1
0
void convertProtoToLuaV2(const caffe::NetParameter &netparam, const char* lua_name, const char* cuda_package)
{
  PACKAGE_TYPE cuda_package_type = CCN2;
  if(std::string(cuda_package) == "ccn2")
    cuda_package_type = CCN2;
  else if(std::string(cuda_package) == "nn")
    cuda_package_type = NN;
  else if(std::string(cuda_package) == "cudnn")
    cuda_package_type = CUDNN;

  std::ofstream ofs (lua_name);

  ofs << "require '" << cuda_package << "'\n";
  ofs << "require 'cunn'\n";
  ofs << "local model = {}\n";
  if(std::string(cuda_package)=="ccn2")
    ofs<< "table.insert(model, {'torch_transpose_dwhb', nn.Transpose({1,4},{1,3},{1,2})})\n";
  else if(std::string(cuda_package)=="nn" || std::string(cuda_package)=="cudnn")
    ofs<< "require 'inn'\n";

  int num_output = netparam.input_shape_size() * 4;

  for (int i=0; i<netparam.layer_size(); ++i)
  {
    std::vector<std::pair<std::string, std::string>> lines;
    auto& layer = netparam.layer(i);
    if(layer.type() == "Convolution")
    {
      auto &param = layer.convolution_param();
      int groups = param.group() == 0 ? 1 : param.group();
      int nInputPlane = layer.blobs(0).shape().dim(1)*groups;
      int nOutputPlane = layer.blobs(0).shape().dim(0);
      //int nOutputPlane = param.num_output();
      num_output = nOutputPlane;
      int kW = param.kernel_w();
      int kH = param.kernel_h();
      int dW = param.stride_w();
      int dH = param.stride_h();
      if(kW==0 || kH==0)
      {
        kW = param.kernel_size();
        kH = kW;
      }
      if(dW==0 || dH==0)
      {
        dW = param.stride();
        dH = dW;
      }
      int pad_w = param.pad_w();
      int pad_h = param.pad_h();
      if(pad_w==0 || pad_h==0)
      {
        pad_w = param.pad();
        pad_h = pad_w;
      }
      if(cuda_package_type == CCN2)
      {
        if(kW != kH || dW != dH || pad_w != pad_h)
        {
          std::cout << "ccn2 only supports square images!\n";
          break;
        }
        char buf[1024];
        sprintf(buf, "ccn2.SpatialConvolution(%d, %d, %d, %d, %d, %d)",
            nInputPlane, nOutputPlane, kW, dW, pad_w, groups);
        lines.emplace_back(layer.name(), buf);
      }
      else if(cuda_package_type == NN)
      {
        if(groups != 1)
        {
          std::cout << "nn supports no groups!\n";
          break;
        }
        char buf[1024];
        sprintf(buf, "nn.SpatialConvolutionMM(%d, %d, %d, %d, %d, %d, %d, %d)",
            nInputPlane, nOutputPlane, kW, kH, dW, dH, pad_w, pad_h);
        lines.emplace_back(layer.name(), buf);
      }
      else
      {
        char buf[1024];
        sprintf(buf, "cudnn.SpatialConvolution(%d, %d, %d, %d, %d, %d, %d, %d, %d)",
            nInputPlane, nOutputPlane, kW, kH, dW, dH, pad_w, pad_h, groups);
        lines.emplace_back(layer.name(), buf);
      }
    }
    if(layer.type() == "Pooling")
    {
      auto &param = layer.pooling_param();
      std::string ptype = param.pool() == caffe::PoolingParameter::MAX ? "Max" : "Avg";
      int kW = param.kernel_w();
      int kH = param.kernel_h();
      int dW = param.stride_w();
      int dH = param.stride_h();
      int padW = param.pad_w();
      int padH = param.pad_h();
      if(kW==0 || kH==0)
      {
        kW = param.kernel_size();
        kH = kW;
      }
      if(dW==0 || dH==0)
      {
        dW = param.stride();
        dH = dW;
      }

      char buf[1024];
      switch(cuda_package_type)
      {
        case CCN2: // ceil mode by default
          if(param.pool() == caffe::PoolingParameter::MAX)
            sprintf(buf, "ccn2.SpatialMaxPooling(%d, %d)", kW, dW);
          else if(param.pool() == caffe::PoolingParameter::AVE)
            sprintf(buf, "ccn2.SpatialAvgPooling(%d, %d)", kW, dW);
          else if(param.pool() == caffe::PoolingParameter::STOCHASTIC)
            THError("Stochastic pooling is not implemented in DHWB format");
          break;
        case CUDNN:
          if(param.pool() == caffe::PoolingParameter::MAX)
            sprintf(buf, "cudnn.SpatialMaxPooling(%d, %d, %d, %d, %d, %d):ceil()", kW, kH, dW, dH, padW, padH);
          else if(param.pool() == caffe::PoolingParameter::AVE)
            sprintf(buf, "cudnn.SpatialAveragePooling(%d, %d, %d, %d, %d, %d):ceil()", kW, kH, dW, dH, padW, padH);
          else if(param.pool() == caffe::PoolingParameter::STOCHASTIC)
            sprintf(buf, "inn.SpatialStochasticPooling(%d, %d, %d, %d)", kW, kH, dW, dH);
          break;
        case NN:
          if(param.pool() == caffe::PoolingParameter::MAX)
            sprintf(buf, "nn.SpatialMaxPooling(%d, %d, %d, %d, %d, %d):ceil()", kW, kH, dW, dH, padW, padH);
          else if(param.pool() == caffe::PoolingParameter::AVE)
            sprintf(buf, "inn.SpatialAveragePooling(%d, %d, %d, %d)", kW, kH, dW, dH); // padding is not supported yet
          else if(param.pool() == caffe::PoolingParameter::STOCHASTIC)
            sprintf(buf, "inn.SpatialStochasticPooling(%d, %d, %d, %d)", kW, kH, dW, dH);
          break;
      }
      lines.emplace_back(layer.name(), buf);
    }
    if(layer.type() == "ReLU")
    {
      if(cuda_package_type == CUDNN)
	lines.emplace_back(layer.name(), "cudnn.ReLU(true)");
      else
        lines.emplace_back(layer.name(), "nn.ReLU(true)");
    }
    if(layer.type() == "Sigmoid")
    {
      if(cuda_package_type == CUDNN)
	lines.emplace_back(layer.name(), "cudnn.Sigmoid(true)");
      else
        lines.emplace_back(layer.name(), "nn.Sigmoid()");
    }
    if(layer.type() == "Tanh")
    {
      if(cuda_package_type == CUDNN)
	lines.emplace_back(layer.name(), "cudnn.Tanh(true)");
      else
        lines.emplace_back(layer.name(), "nn.Tanh()");
    }
    if(layer.type() == "LRN")
    {
      auto &param = layer.lrn_param();
      int local_size = param.local_size();
      float alpha = param.alpha();
      float beta = param.beta();
      float k = param.k();
      char buf[1024];
      if(std::string(cuda_package) == "ccn2")
        sprintf(buf, "ccn2.SpatialCrossResponseNormalization(%d, %.6f, %.4f, %f)", local_size, alpha, beta, k);
      else
        sprintf(buf, "inn.SpatialCrossResponseNormalization(%d, %.6f, %.4f, %f)", local_size, alpha, beta, k);
      lines.emplace_back(layer.name(), buf);
    }
    if(layer.type() == "InnerProduct")
    {
      auto &param = layer.inner_product_param();
      int nInputPlane = layer.blobs(0).shape().dim(1);
      int nOutputPlane = param.num_output();
      char buf[1024];
      sprintf(buf, "nn.Linear(%d, %d)", nInputPlane, nOutputPlane);
      if(num_output != nInputPlane)
      {
        if(std::string(cuda_package) == "ccn2")
          lines.emplace_back("torch_transpose_bdwh", "nn.Transpose({4,1},{4,2},{4,3})");
        lines.emplace_back("torch_view", "nn.View(-1):setNumInputDims(3)");
      }
      lines.emplace_back(layer.name(), buf);
      num_output = nOutputPlane;
    }
    if(layer.type() == "Dropout")
    {
      char buf[1024];
      sprintf(buf, "nn.Dropout(%f)", layer.dropout_param().dropout_ratio());
      lines.emplace_back(layer.name(), buf);
    }
    if(layer.type()=="SoftmaxWithLoss")
    {
      lines.emplace_back(layer.name(), "nn.SoftMax()");
    }
    if(layer.type()=="Softmax")
    {
      lines.emplace_back(layer.name(), "nn.SoftMax()");
    }

    if(!lines.empty())
      for(auto& it: lines)
        ofs << "table.insert(model, {'" << it.first << "', " << it.second << "})\n";
    else
    {
      ofs << "-- warning: module '" << layer.name() << "' not found\n";
      std::cout << "warning: module '" << layer.name() << "' not found\n";
    }
  }
  ofs << "return model";
}