support signal handler

microsoft · Feb 26, 2016 · 475ab87 · 475ab87
1 parent 682cef8
commit 475ab87
Show file tree

Hide file tree

Showing 4 changed files with 28 additions and 7 deletions.
diff --git a/src/caffe/util/signal_handler.cpp b/src/caffe/util/signal_handler.cpp
@@ -13,9 +13,15 @@ namespace {
 
   void handle_signal(int signal) {
     switch (signal) {
+#ifdef _MSC_VER
+    case SIGBREAK://there is no SIGHUP in windows, take SIGBREAK instead.
+      got_sighup = true;
+      break;
+#else
     case SIGHUP:
       got_sighup = true;
       break;
+#endif
     case SIGINT:
       got_sigint = true;
       break;
@@ -27,7 +33,14 @@ namespace {
       LOG(FATAL) << "Tried to hookup signal handlers more than once.";
     }
     already_hooked_up = true;
-
+#ifdef _MSC_VER
+    if (signal(SIGBREAK, handle_signal) == SIG_ERR) {
+      LOG(FATAL) << "Cannot install SIGBREAK handler.";
+    }
+    if (signal(SIGINT, handle_signal) == SIG_ERR) {
+      LOG(FATAL) << "Cannot install SIGINT handler.";
+    }
+#else
     struct sigaction sa;
     // Setup the handler
     sa.sa_handler = &handle_signal;
@@ -42,11 +55,20 @@ namespace {
     if (sigaction(SIGINT, &sa, NULL) == -1) {
       LOG(FATAL) << "Cannot install SIGINT handler.";
     }
+#endif
   }
 
   // Set the signal handlers to the default.
   void UnhookHandler() {
     if (already_hooked_up) {
+#ifdef _MSC_VER
+      if (signal(SIGBREAK, SIG_DFL) == SIG_ERR) {
+        LOG(FATAL) << "Cannot uninstall SIGHUP handler.";
+      }
+      if (signal(SIGINT, SIG_DFL) == SIG_ERR) {
+        LOG(FATAL) << "Cannot uninstall SIGINT handler.";
+      }
+#else
       struct sigaction sa;
       // Setup the sighub handler
       sa.sa_handler = SIG_DFL;
@@ -61,7 +83,7 @@ namespace {
       if (sigaction(SIGINT, &sa, NULL) == -1) {
         LOG(FATAL) << "Cannot uninstall SIGINT handler.";
       }
-
+#endif
       already_hooked_up = false;
     }
   }

diff --git a/tools/caffe.cpp b/tools/caffe.cpp
@@ -197,19 +197,14 @@ int train() {
     Caffe::set_solver_count(gpus.size());
   }
 
-#if !defined(_MSC_VER)
-  // Signals are not properly supported in Windows.
   caffe::SignalHandler signal_handler(
         GetRequestedAction(FLAGS_sigint_effect),
         GetRequestedAction(FLAGS_sighup_effect));
-#endif
 
   shared_ptr<caffe::Solver<float> >
       solver(caffe::SolverRegistry<float>::CreateSolver(solver_param));
 
-#if !defined(_MSC_VER)
   solver->SetActionFunction(signal_handler.GetActionFunction());
-#endif
 
   if (FLAGS_snapshot.size()) {
     LOG(INFO) << "Resuming from " << FLAGS_snapshot;

diff --git a/windows/libcaffe/libcaffe.vcxproj b/windows/libcaffe/libcaffe.vcxproj
@@ -186,6 +186,7 @@
     <ClCompile Include="..\..\src\caffe\util\insert_splits.cpp" />
     <ClCompile Include="..\..\src\caffe\util\io.cpp" />
     <ClCompile Include="..\..\src\caffe\util\math_functions.cpp" />
+    <ClCompile Include="..\..\src\caffe\util\signal_handler.cpp" />
     <ClCompile Include="..\..\src\caffe\util\upgrade_proto.cpp" />
   </ItemGroup>
   <ItemGroup>

diff --git a/windows/libcaffe/libcaffe.vcxproj.filters b/windows/libcaffe/libcaffe.vcxproj.filters
@@ -330,6 +330,9 @@
     <ClCompile Include="..\..\src\caffe\layers\mil_layer.cpp">
       <Filter>src\layers</Filter>
     </ClCompile>
+    <ClCompile Include="..\..\src\caffe\util\signal_handler.cpp">
+      <Filter>src\util</Filter>
+    </ClCompile>
   </ItemGroup>
   <ItemGroup>
     <ClInclude Include="..\..\include\caffe\proto\caffe.pb.h">