diff --git a/include/caffe/util/io.hpp b/include/caffe/util/io.hpp index 6b7332548b0..1a599883ca3 100644 --- a/include/caffe/util/io.hpp +++ b/include/caffe/util/io.hpp @@ -2,33 +2,51 @@ #define CAFFE_UTIL_IO_H_ #include +#include +#include // NOLINT(readability/streams) #include #include "google/protobuf/message.h" #include "caffe/common.hpp" #include "caffe/proto/caffe.pb.h" +#include "caffe/util/format.hpp" + +#ifndef CAFFE_TMP_DIR_RETRIES +#define CAFFE_TMP_DIR_RETRIES 100 +#endif namespace caffe { using ::google::protobuf::Message; using ::boost::filesystem::path; -inline void MakeTempFilename(string* temp_filename) { - temp_filename->clear(); - const path& model = boost::filesystem::temp_directory_path() - /"caffe_test.%%%%%%"; - *temp_filename = boost::filesystem::unique_path(model).string(); -} - inline void MakeTempDir(string* temp_dirname) { temp_dirname->clear(); - const path& model = boost::filesystem::temp_directory_path() - /"caffe_test.%%%%%%"; - const path& dir = boost::filesystem::unique_path(model).string(); - bool directoryCreated = boost::filesystem::create_directory(dir); - CHECK(directoryCreated); - *temp_dirname = dir.string(); + const path& model = + boost::filesystem::temp_directory_path()/"caffe_test.%%%%-%%%%"; + for ( int i = 0; i < CAFFE_TMP_DIR_RETRIES; i++ ) { + const path& dir = boost::filesystem::unique_path(model).string(); + bool done = boost::filesystem::create_directory(dir); + if ( done ) { + *temp_dirname = dir.string(); + return; + } + } + LOG(FATAL) << "Failed to create a temporary directory."; +} + +inline void MakeTempFilename(string* temp_filename) { + static path temp_files_subpath; + static uint64_t next_temp_file = 0; + temp_filename->clear(); + if ( temp_files_subpath.empty() ) { + string path_string=""; + MakeTempDir(&path_string); + temp_files_subpath = path_string; + } + *temp_filename = + (temp_files_subpath/caffe::format_int(next_temp_file++, 9)).string(); } bool ReadProtoFromTextFile(const char* filename, Message* proto);