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 [resume_point_file]
   7
   8 #include <cuda_runtime.h>
   9
  10 #include <cstring>
  11
  12 #include "caffe/caffe.hpp"
  13
  14 using namespace caffe;
  15
  16 int main(int argc, char** argv) {
  17   ::google::InitGoogleLogging(argv[0]);
  18   if (argc < 2) {
  19     LOG(ERROR) << "Usage: train_net solver_proto_file [resume_point_file]";
  20     return 0;
  21   }
  22
  23   SolverParameter solver_param;
  24   ReadProtoFromTextFile(argv[1], &solver_param);
  25
  26   LOG(INFO) << "Starting Optimization";
  27   SGDSolver<float> solver(solver_param);
  28   if (argc == 3) {
  29     LOG(INFO) << "Resuming from " << argv[2];
  30     solver.Solve(argv[2]);
  31   } else {
  32     solver.Solve();
  33   }
  34   LOG(INFO) << "Optimization Done.";
  35
  36   return 0;
  37 }