caffe2 first sight

#include 
// Caffe2
#include 
#include "caffe2/core/context.h"
#include "caffe2/core/operator.h"
#include "caffe2/core/predictor.h"
#include "caffe2/core/tensor.h"
#include "caffe2/utils/math.h"
#include 
#include 
#include 
#include 
// OpenCV
#include 
#include 
#include 
#include 
#include 
void setgpu(caffe2::NetDef& net, int id){
    net.mutable_device_option()->set_device_type(1);
    net.mutable_device_option()->set_cuda_gpu_id(id);
    for(int i = 0; i < net.op_size(); ++i){
        net.mutable_op(i)->mutable_device_option()->set_device_type(1);
        net.mutable_op(i)->mutable_device_option()->set_cuda_gpu_id(0);
    }
}
int main(int argc, char* argv[])
{
    std::cout <<"1" <<std::endl;
    //typedef caffe2::TensorCPU tensorType;
    typedef caffe2::TensorCUDA TensorType;
    // std::string image_filename = us::any_cast(parsedArgs["image"]);
    // cv::Mat img = cv::imread(image_filename, CV_LOAD_IMAGE_COLOR);
    // cv::Mat img = cv::Mat(64,64,CV_32FC3,128);
    // cv::namedWindow( "Display window", cv::WINDOW_AUTOSIZE );// Create a window for display.
    // cv::imshow( "Display w30indow", img2 ); // Show our image inside it.
    // cv::waitKey(0); // Wait for a keystroke in the window
    caffe2::NetDef init_net, predict_net;

    CAFFE_ENFORCE(ReadProtoFromFile("/home/amax/Yang/caffe2/models/init_net.pb", &init_net));
    setgpu(init_net, 1);
    std::cout <<"2" <<std::endl;
    CAFFE_ENFORCE(ReadProtoFromFile("/home/amax/Yang/caffe2/models/predict_net.pb", &predict_net));
    setgpu(predict_net, 1);
    std::cout <<"3" <<std::endl;
    std::vector outputVec;
    int w,h,c; w=224; h=224; c=3;
    TensorType input(std::vector<int>({1, c, h, w}));
     TRYING TO ACCESS TENSOR DATA (RESULTS IN ERROR FOR caffe2::TensorCUDA)
    // caffe2::CPUContext cpu_context;
    // caffe2::TensorCPU cpu_input(input, &cpu_context);
    // float* input_data = cpu_input.template mutable_data();
    float* input_data = input.mutable_data<float>();
    // create input image (simply a square of size 10x10 in a 64x64 image)
    // test image for segmentation
    /*int counter = 0;
    for (int x = 0; x < h; x++)
        for (int y = 0; y < w; y++)
        {
            for (int j = 0; j < c; j++)
            {
                input_data[counter] = 0;
                if (x>20 && x<30 && y>20 && y<30)
                    input_data[counter] = 1;
                counter++;
        }
    }*/
    std::cout <<"4" <<std::endl;
    // show input image
    cv::Mat img = cv::Mat(224,224,CV_32FC3);
    cv::namedWindow( "Display window", cv::WINDOW_AUTOSIZE );// Create a window for display.
    cv::imshow( "Display window", img ); // Show our image inside it.
    cv::waitKey(0); // Wait for a keystroke in the window
    // run network
    clock_t t1,t2;
    std::vector inputVec{input};
    caffe2::Workspace ws;
    t1=clock();
    ws.RunNetOnce(init_net);
    t2=clock();
    std::cout <<(double)(t2-t1) / CLOCKS_PER_SEC <<std::endl;
    ws.CreateNet(predict_net, true);
    std::cout << inputVec.size() << std::endl;
    for (auto i = 0; i < inputVec.size(); ++i)
    {
        auto* blob = ws.GetBlob(predict_net.external_input(i));
        TensorType* tensor = blob->template GetMutable();
        tensor->ResizeLike(inputVec[i]);
        tensor->ShareData(inputVec[i]);
    }
    std::cout <<"5" <<std::endl;

    t1=clock();
    CAFFE_ENFORCE(ws.RunNet(predict_net.name()));
    t2=clock();
    std::cout <<(double)(t2-t1) / CLOCKS_PER_SEC <<std::endl;

    t1=clock();
    for(int i = 0; i < 1000; ++i)
        ws.RunNet(predict_net.name());
    t2=clock();
    std::cout <<(double)(t2-t1) / CLOCKS_PER_SEC/1000 <<std::endl;

    t1=clock();
    ws.RunNet(predict_net.name());
    t2=clock();
    std::cout <<(double)(t2-t1) / CLOCKS_PER_SEC <<std::endl;

    std::cout <<"6" <<std::endl;
    return 1;
}

你可能感兴趣的:(caffe2)