resnet-multi-gpu.cc

#include "NexusCpp.hpp"

// 一些全局变量
constexpr int N = 16;       // batch size
constexpr int C = 1;        // channels 通道数
constexpr int H = 224;      // 图像的高
constexpr int W = 224;      // 图像的宽
constexpr int NCLASS = 10;  // 分类的总类别数

auto DATASHAPE = TShape({N, C, H, W});  // 输入图像的形状
auto LABELSHAPE = TShape({N, NCLASS});  // 标签的形状
auto PREDSHAPE = LABELSHAPE;  // 预测的形状与标签的形状相同

Symbol ResNet(int num_class) {
  int inchannel = 64, channel = 0;
  int stride = 1;
  int l = 1;
  auto data = Symbol::Variable("data");
  auto conv1 = Convolution("conv1", data, 3, 64, true, 1, 1, 1);
  auto bat1 = BatchNormalization("bat1", conv1);
  auto relu1 = Relu("relu1", bat1);
  auto x = relu1;
  for (int i = 1; i <= 4; i++) {
    if (i != 1) stride = 2;
    if (i == 1)
      channel = 64;
    else if (i == 2)
      channel = 128;
    else if (i == 3)
      channel = 256;
    else
      channel = 512;
    for (int j = 1; j <= 2; j++) {
      auto downsample = x;
      auto conv2 = Convolution("conv" + std::to_string(++l), x, 3, channel,
                               true, stride, 1, 1);
      auto bat2 = BatchNormalization("bat" + std::to_string(l), conv2);
      auto relu2 = Relu("relu" + std::to_string(l), bat2);
      if (stride != 1 || inchannel != channel) {
        downsample = Convolution("downsample" + std::to_string(l), x, 1,
                                 channel, true, stride);
        downsample = BatchNormalization("downsample_batch" + std::to_string(l),
                                        downsample);
      }
      inchannel = channel;
      auto conv3 = Convolution("conv" + std::to_string(++l), relu2, 3, channel,
                               true, 1, 1, 1);
      auto bat3 = BatchNormalization("bat" + std::to_string(l), conv3);
      auto out = Sum(std::string("out") + std::to_string(l), 2,
                     {bat3, downsample});  //此处需x要调试
      x = Relu("x" + std::to_string(l), out);
    }
  }
  x = Pooling("pool1", x, 4, 4);
  auto fc4 = FullyConnected("fc4", x, num_class);
  auto label = Symbol::Variable("label");
  auto loss = LogSoftmaxLoss("softmax", fc4, label);
  return loss;
}

// 解析命令行参数
cxxopts::ParseResult parse_commandline(int argc, char** argv) {
  cxxopts::Options options(argv[0], "DenseNet 训练程序");
  // clang-format off
    options.add_options()
    ("l,lr", "学习率", cxxopts::value<float>()->default_value("1e-4"))
    ("g,gpu", "GPU ID", cxxopts::value<int>()->default_value("0"))
    ("e,epochs", "训练轮数", cxxopts::value<int>()->default_value("5"))
    ("o,output", "保存模型名称", cxxopts::value<std::string>()->default_value(argv[0]))
    ("h,help", "打印帮助");
  // clang-format on

  auto result = options.parse(argc, argv);

  if (result.count("help")) {
    std::cout << options.help() << std::endl;
    exit(0);
  }

  return result;
}

int main(int argc, char** argv) {
  auto parsed_options = parse_commandline(argc, argv);
  auto nn = ResNet(NCLASS);

  auto gpu0 = Context::GPU(0);
  auto gpu1 = Context::GPU(1);
  auto gpu2 = Context::GPU(2);
  auto gpu3 = Context::GPU(3);
  auto cpu = Context::CPU();

  std::map<std::string, Tensor> args_map;
  args_map["data"] = Tensor(DATASHAPE, kFloat32, gpu0);
  args_map["label"] = Tensor(LABELSHAPE, kFloat32, gpu3);

  nn.InferArgsMap(gpu0, &args_map, args_map, true);

  // ! 非常重要，必须放前面，因为后面 PartialBind 会改写 nn（把 nn 进行分割），
  // ! 所以如果在后面调用 nn.ListArguments()，就会列出不完整的 arguments
  auto arguments = nn.ListArguments(); 

  auto exec0 = nn.PartialBind(gpu0, args_map, "data", "out5");
  auto exec1 = nn.PartialBind(gpu1, args_map, "x5", "out11");
  auto exec2 = nn.PartialBind(gpu2, args_map, "x11", "fc4");
  auto exec3 = nn.PartialBind(gpu3, args_map, "softmax", "softmax");

  Engine::Get()->WaitForAll();
  
  std::vector<std::pair<std::string, std::string>> dataset_params = {
      {"batch_size", std::to_string(N)},
      {"shuffle", "true"},
      {"input_height", std::to_string(H)},
      {"input_width", std::to_string(W)}};

  std::vector<std::pair<std::string, std::string>> testset_params = {
      {"image", "/data/mnist/t10k-images-idx3-ubyte"},
      {"label", "/data/mnist/t10k-labels-idx1-ubyte"},
      {"batch_size", std::to_string(N)},
      {"shuffle", "true"},
      {"input_height", std::to_string(H)},
      {"input_width", std::to_string(W)}};

  auto dataloader = std::make_shared<nexus::io::MNISTLoader>();
  dataloader->Init(dataset_params);

  auto testloader = std::make_shared<nexus::io::MNISTLoader>();
  testloader->Init(testset_params);

  float learning_rate = parsed_options["lr"].as<float>();
  float weight_decay = 1e-5;
  Optimizer* opt0 = OptimizerRegistry::Find("sgd", gpu0);
  opt0->SetParam("rescale_grad", 1.0 / N)
      ->SetParam("lr", learning_rate)
      ->SetParam("wd", weight_decay);
  Optimizer* opt1 = OptimizerRegistry::Find("sgd", gpu1);
  opt1->SetParam("rescale_grad", 1.0 / N)
      ->SetParam("lr", learning_rate)
      ->SetParam("wd", weight_decay);
  Optimizer* opt2 = OptimizerRegistry::Find("sgd", gpu2);
  opt2->SetParam("rescale_grad", 1.0 / N)
      ->SetParam("lr", learning_rate)
      ->SetParam("wd", weight_decay);
  Optimizer* opt3 = OptimizerRegistry::Find("sgd", gpu3);
  opt3->SetParam("rescale_grad", 1.0 / N)
      ->SetParam("lr", learning_rate)
      ->SetParam("wd", weight_decay);

  testloader->BeforeFirst();
  int count = 0;
  for (int iter = 1; iter <= parsed_options["epochs"].as<int>(); ++iter) {
    dataloader->BeforeFirst();
    while (dataloader->Next()) {
      io::TensorBatch batch = dataloader->Value();
      Tensor labels(LABELSHAPE, kFloat32, cpu);
      LabelOneHot(labels, batch.labels);

      // 前向计算
      batch.data[0] >> exec0->GetArgByName("data");
      exec0->Forward(true);
      exec0->GetForwardOutputs()[0] >> exec1->GetArgByName("x5_input");
      exec1->Forward(true);
      exec1->GetForwardOutputs()[0] >> exec2->GetArgByName("x11_input");
      exec2->Forward(true);
      exec2->GetForwardOutputs()[0] >> exec3->GetArgByName("softmax_input");
      labels >> exec3->GetArgByName("label");
      exec3->Forward(true);

      if (count % 100 == 0) {
        auto loss = exec3->GetForwardOutputs()[0].To(cpu);
        float loss_mean = 0.0;
        for (int i = 0; i < N; i++) loss_mean += loss[i].Item();
        loss_mean /= N;
        printf("count = %d, loss = %.3f, ", count, loss_mean);
        
        auto pred = exec2->GetForwardOutputs()[0].To(cpu);
        auto acc = Acc(pred, batch.labels);
        printf("acc = %.2f%%\n", acc * 100.0);
      }

      // 每 1000 个 batch，学习率减半（针对 SGD）
      // if ((count + 1) % 1000 == 0) {
      //   learning_rate *= 0.5;
      //   opt->SetParam("lr", learning_rate);
      // }

      // 反向传播
      exec3->Backward();
      exec2->Backward({exec3->GetGradByName("softmax_input")});
      exec1->Backward({exec2->GetGradByName("x11_input")});
      exec0->Backward({exec1->GetGradByName("x5_input")});

      // 使用优化器来优化网络参数
      int i = 0;
      for (auto& name : arguments) {
        if (name == "data" || name == "label") continue;
        if (exec0->HasGradByName(name)) {
          opt0->Update(i, exec0->GetArgByName(name), exec0->GetGradByName(name), gpu0);
        } else if (exec1->HasArgByName(name)) {
          opt1->Update(i, exec1->GetArgByName(name), exec1->GetGradByName(name), gpu1);
        } else if (exec2->HasArgByName(name)) {
          opt2->Update(i, exec2->GetArgByName(name), exec2->GetGradByName(name), gpu2);
        } else if (exec3->HasArgByName(name)) {
          opt3->Update(i, exec3->GetArgByName(name), exec3->GetGradByName(name), gpu3);
        } else {
          LOG(FATAL) << "Unknown arg name!\n";
        }
        i++;
      }
      count++;
    }
  }

  // 重要：等待引擎上所有计算完成
  Engine::Get()->WaitForAll();

  // 通知引擎关闭
  Engine::Get()->NotifyShutdown();

  return 0;
}