finetune code
authorYangqing Jia <jiayq84@gmail.com>
Fri, 15 Nov 2013 05:12:23 +0000 (21:12 -0800)
committerYangqing Jia <jiayq84@gmail.com>
Fri, 15 Nov 2013 05:12:23 +0000 (21:12 -0800)
examples/finetune_net.cpp [new file with mode: 0644]

diff --git a/examples/finetune_net.cpp b/examples/finetune_net.cpp
new file mode 100644 (file)
index 0000000..d35b425
--- /dev/null
@@ -0,0 +1,36 @@
+// Copyright 2013 Yangqing Jia
+//
+// This is a simple script that allows one to quickly finetune a network.
+// Usage:
+//    finetune_net solver_proto_file pretrained_net
+
+#include <cuda_runtime.h>
+
+#include <cstring>
+
+#include "caffe/caffe.hpp"
+
+using namespace caffe;
+
+int main(int argc, char** argv) {
+  ::google::InitGoogleLogging(argv[0]);
+  if (argc < 2) {
+    LOG(ERROR) << "Usage: finetune_net solver_proto_file pretrained_net";
+    return 0;
+  }
+
+  Caffe::SetDevice(0);
+  Caffe::set_mode(Caffe::GPU);
+
+  SolverParameter solver_param;
+  ReadProtoFromTextFile(argv[1], &solver_param);
+
+  LOG(INFO) << "Starting Optimization";
+  SGDSolver<float> solver(solver_param);
+  LOG(INFO) << "Loading from " << argv[2];
+  solver.net()->CopyTrainedLayersFrom(string(argv[2]));
+  solver.Solve();
+  LOG(INFO) << "Optimization Done.";
+
+  return 0;
+}