Caffe2 - C++ API
A deep learning, cross platform ML framework
data_channel_mpi_smoke.cpp
1 #include <THD/base/data_channels/DataChannelMPI.hpp>
2 
3 #include <unistd.h>
4 #include <cassert>
5 #include <iostream>
6 #include <memory>
7 
8 constexpr int WORKERS_NUM = 2;
9 
10 int main(int argc, char** argv) {
11  if (argc == 1) {
12  execlp(
13  "mpirun",
14  "mpirun",
15  "-n",
16  std::to_string(WORKERS_NUM + 1).data(),
17  argv[0],
18  "1",
19  NULL);
20  }
21 
22  auto dataChannel = std::make_shared<thd::DataChannelMPI>();
23  assert(dataChannel->init());
24  assert(dataChannel->getNumProcesses() == (WORKERS_NUM + 1));
25  std::cout << "OK (id: " << dataChannel->getRank() << ")" << std::endl;
26  return 0;
27 }