examples/finetune_net.cpp

   1 // Copyright 2013 Yangqing Jia
   2 //
   3 // This is a simple script that allows one to quickly finetune a network.
   4 // Usage:
   5 //    finetune_net solver_proto_file pretrained_net
   6
   7 #include <cuda_runtime.h>
   8
   9 #include <cstring>
  10
  11 #include "caffe/caffe.hpp"
  12
  13 using namespace caffe;
  14
  15 int main(int argc, char** argv) {
  16   ::google::InitGoogleLogging(argv[0]);
  17   if (argc < 2) {
  18     LOG(ERROR) << "Usage: finetune_net solver_proto_file pretrained_net";
  19     return 0;
  20   }
  21
  22   SolverParameter solver_param;
  23   ReadProtoFromTextFile(argv[1], &solver_param);
  24
  25   LOG(INFO) << "Starting Optimization";
  26   SGDSolver<float> solver(solver_param);
  27   LOG(INFO) << "Loading from " << argv[2];
  28   solver.net()->CopyTrainedLayersFrom(string(argv[2]));
  29   solver.Solve();
  30   LOG(INFO) << "Optimization Done.";
  31
  32   return 0;
  33 }