Caffe2 - C++ API
A deep learning, cross platform ML framework
run_plan_mpi.cc
1 
17 #include <mpi.h>
18 
19 #include "caffe2/core/init.h"
20 #include "caffe2/core/operator.h"
21 #include "caffe2/proto/caffe2.pb.h"
22 #include "caffe2/utils/proto_utils.h"
23 #include "caffe2/core/logging.h"
24 
25 CAFFE2_DEFINE_string(plan, "", "The given path to the plan protobuffer.");
26 
27 int main(int argc, char** argv) {
28  caffe2::SetUsageMessage("Runs a caffe2 plan that has MPI operators in it.");
29  int mpi_ret;
30  MPI_Init_thread(&argc, &argv, MPI_THREAD_MULTIPLE, &mpi_ret);
31  if (mpi_ret != MPI_THREAD_MULTIPLE &&
32  mpi_ret != MPI_THREAD_SERIALIZED) {
33  std::cerr << "Caffe2 MPI requires the underlying MPI to support the "
34  "MPI_THREAD_SERIALIZED or MPI_THREAD_MULTIPLE mode.\n";
35  return 1;
36  }
37  caffe2::GlobalInit(&argc, &argv);
38  LOG(INFO) << "Loading plan: " << caffe2::FLAGS_plan;
39  caffe2::PlanDef plan_def;
40  CAFFE_ENFORCE(ReadProtoFromFile(caffe2::FLAGS_plan, &plan_def));
41  std::unique_ptr<caffe2::Workspace> workspace(new caffe2::Workspace());
42  workspace->RunPlan(plan_def);
43 
44  // This is to allow us to use memory leak checks.
45  google::protobuf::ShutdownProtobufLibrary();
46  MPI_Finalize();
47  return 0;
48 }
bool GlobalInit(int *pargc, char ***pargv)
Initialize the global environment of caffe2.
Definition: init.cc:34
Workspace is a class that holds all the related objects created during runtime: (1) all blobs...
Definition: workspace.h:63
void SetUsageMessage(const string &str)
Sets the usage message when a commandline tool is called with "--help".
Definition: flags.cc:68