examples/train_net.cpp

   1 // Copyright 2013 Yangqing Jia
   2 //
   3 // This is a simple script that allows one to quickly train a network whose
   4 // parameters are specified by text format protocol buffers.
   5 // Usage:
   6 //    train_net net_proto_file solver_proto_file
   7
   8 #include <cuda_runtime.h>
   9
  10 #include <cstring>
  11
  12 #include "caffe/blob.hpp"
  13 #include "caffe/common.hpp"
  14 #include "caffe/net.hpp"
  15 #include "caffe/filler.hpp"
  16 #include "caffe/proto/caffe.pb.h"
  17 #include "caffe/util/io.hpp"
  18 #include "caffe/solver.hpp"
  19
  20 using namespace caffe;
  21
  22 int main(int argc, char** argv) {
  23   cudaSetDevice(0);
  24   Caffe::set_mode(Caffe::GPU);
  25   Caffe::set_phase(Caffe::TRAIN);
  26
  27   NetParameter net_param;
  28   ReadProtoFromTextFile(argv[1], &net_param);
  29   vector<Blob<float>*> bottom_vec;
  30   Net<float> caffe_net(net_param, bottom_vec);
  31
  32   SolverParameter solver_param;
  33   ReadProtoFromTextFile(argv[2], &solver_param);
  34
  35   LOG(ERROR) << "Starting Optimization";
  36   SGDSolver<float> solver(solver_param);
  37   solver.Solve(&caffe_net);
  38   LOG(ERROR) << "Optimization Done.";
  39
  40   return 0;
  41 }