|
2 | 2 | #include <google/protobuf/io/zero_copy_stream_impl.h> |
3 | 3 | #include <google/protobuf/text_format.h> |
4 | 4 |
|
| 5 | +#include <boost/filesystem.hpp> |
| 6 | + |
5 | 7 | #include <map> |
6 | 8 | #include <string> |
7 | 9 |
|
@@ -1101,12 +1103,31 @@ bool UpgradeSolverAsNeeded(const string& param_file, SolverParameter* param) { |
1101 | 1103 | return success; |
1102 | 1104 | } |
1103 | 1105 |
|
| 1106 | +// Replaces snapshot_prefix of SolverParameter if it is not specified |
| 1107 | +// or is set to directory |
| 1108 | +void UpgradeSnapshotPrefixProperty(const string& param_file, |
| 1109 | + SolverParameter* param) { |
| 1110 | + using boost::filesystem::path; |
| 1111 | + using boost::filesystem::is_directory; |
| 1112 | + if (!param->has_snapshot_prefix()) { |
| 1113 | + param->set_snapshot_prefix(path(param_file).replace_extension().string()); |
| 1114 | + LOG(INFO) << "snapshot_prefix was not specified and is set to " |
| 1115 | + + param->snapshot_prefix(); |
| 1116 | + } else if (is_directory(param->snapshot_prefix())) { |
| 1117 | + param->set_snapshot_prefix((path(param->snapshot_prefix()) / |
| 1118 | + path(param_file).stem()).string()); |
| 1119 | + LOG(INFO) << "snapshot_prefix was a directory and is replaced to " |
| 1120 | + + param->snapshot_prefix(); |
| 1121 | + } |
| 1122 | +} |
| 1123 | + |
1104 | 1124 | // Read parameters from a file into a SolverParameter proto message. |
1105 | 1125 | void ReadSolverParamsFromTextFileOrDie(const string& param_file, |
1106 | 1126 | SolverParameter* param) { |
1107 | 1127 | CHECK(ReadProtoFromTextFile(param_file, param)) |
1108 | 1128 | << "Failed to parse SolverParameter file: " << param_file; |
1109 | 1129 | UpgradeSolverAsNeeded(param_file, param); |
| 1130 | + UpgradeSnapshotPrefixProperty(param_file, param); |
1110 | 1131 | } |
1111 | 1132 |
|
1112 | 1133 | } // namespace caffe |
0 commit comments