forked from BVLC/caffe
/
train_net.cpp
37 lines (29 loc) · 879 Bytes
/
train_net.cpp
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
// Copyright 2013 Yangqing Jia
//
// This is a simple script that allows one to quickly train a network whose
// parameters are specified by text format protocol buffers.
// Usage:
// train_net net_proto_file solver_proto_file [resume_point_file]
#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: train_net solver_proto_file [resume_point_file]";
return 0;
}
SolverParameter solver_param;
ReadProtoFromTextFile(argv[1], &solver_param);
LOG(INFO) << "Starting Optimization";
SGDSolver<float> solver(solver_param);
if (argc == 3) {
LOG(INFO) << "Resuming from " << argv[2];
solver.Solve(argv[2]);
} else {
solver.Solve();
}
LOG(INFO) << "Optimization Done.";
return 0;
}