sane h5df file type check for weights

BVLC · Mar 15, 2017 · 623b515 · 623b515
1 parent 317d162
commit 623b515
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 5 deletions.
diff --git a/Makefile b/Makefile
@@ -178,7 +178,7 @@ ifneq ($(CPU_ONLY), 1)
 	LIBRARIES := cudart cublas curand
 endif
 
-LIBRARIES += glog gflags protobuf boost_system boost_filesystem m hdf5_hl hdf5
+LIBRARIES += glog gflags protobuf boost_system boost_filesystem m hdf5_hl hdf5 hdf5_cpp
 
 # handle IO dependencies
 USE_LEVELDB ?= 1

diff --git a/cmake/Dependencies.cmake b/cmake/Dependencies.cmake
@@ -43,9 +43,9 @@ list(APPEND Caffe_LINKER_LIBS PUBLIC ${GFLAGS_LIBRARIES})
 include(cmake/ProtoBuf.cmake)
 
 # ---[ HDF5
-find_package(HDF5 COMPONENTS HL REQUIRED)
+find_package(HDF5 COMPONENTS HL CXX REQUIRED)
 list(APPEND Caffe_INCLUDE_DIRS PUBLIC ${HDF5_INCLUDE_DIRS})
-list(APPEND Caffe_LINKER_LIBS PUBLIC ${HDF5_LIBRARIES} ${HDF5_HL_LIBRARIES})
+list(APPEND Caffe_LINKER_LIBS PUBLIC ${HDF5_LIBRARIES} ${HDF5_HL_LIBRARIES} ${HDF5_CXX_LIBRARIES})
 
 # ---[ LMDB
 if(USE_LMDB)

diff --git a/src/caffe/net.cpp b/src/caffe/net.cpp
@@ -5,6 +5,7 @@
 #include <utility>
 #include <vector>
 
+#include "H5Cpp.h"
 #include "hdf5.h"
 
 #include "caffe/common.hpp"
@@ -769,8 +770,7 @@ void Net<Dtype>::CopyTrainedLayersFrom(const NetParameter& param) {
 
 template <typename Dtype>
 void Net<Dtype>::CopyTrainedLayersFrom(const string trained_filename) {
-  if (trained_filename.size() >= 3 &&
-      trained_filename.compare(trained_filename.size() - 3, 3, ".h5") == 0) {
+  if (H5::H5File::isHdf5(trained_filename.c_str())) {
     CopyTrainedLayersFromHDF5(trained_filename);
   } else {
     CopyTrainedLayersFromBinaryProto(trained_filename);