mmdeploy/csrc/net/ncnn/ncnn_net.cpp

// Copyright (c) OpenMMLab. All rights reserved.

#include "ncnn_net.h"

#include "core/logger.h"
#include "core/model.h"
#include "core/utils/formatter.h"
#include "ncnn_ops_register.h"

namespace mmdeploy {

NCNNNet::~NCNNNet() {}

Result<void> ncnn_status(int code) {
  if (code == 0) {
    return success();
  }
  return Status(eFail);
}

Result<void> NCNNNet::Init(const Value& args) {
  auto& context = args["context"];
  device_ = context["device"].get<Device>();
  stream_ = context["stream"].get<Stream>();
  if (args.contains("use_vulkan")) {
    net_.opt.use_vulkan_compute = args["use_vulkan"].get<bool>();
  }
  if (!device_.is_host()) {
    return Status(eNotSupported);
  }

  auto name = args["name"].get<std::string>();
  auto model = context["model"].get<Model>();
  OUTCOME_TRY(auto config, model.GetModelConfig(name));
  auto precision = config.precision;
  if (precision == "FP16") {
    net_.opt.use_fp16_packed = true;
    net_.opt.use_fp16_storage = true;
    net_.opt.use_fp16_arithmetic = true;
  } else if (precision == "INT8") {
    // in android platform, ncnn will automatically start FP16 accelerate.
    // In INT8 case, we set fp16 as false explicitly.
    net_.opt.use_int8_packed = true;
    net_.opt.use_int8_storage = true;
    net_.opt.use_int8_arithmetic = true;
    net_.opt.use_fp16_packed = false;
    net_.opt.use_fp16_storage = false;
    net_.opt.use_fp16_arithmetic = false;
  } else {
    // in android platform, ncnn will automatically start FP16 accelerate.
    // In FP32 case, we set fp16 as false explicitly.
    net_.opt.use_fp16_packed = false;
    net_.opt.use_fp16_storage = false;
    net_.opt.use_fp16_arithmetic = false;
  }
  OUTCOME_TRY(params_, model.ReadFile(config.net));
  OUTCOME_TRY(weights_, model.ReadFile(config.weights));
  register_mmdeploy_custom_layers(net_);

  OUTCOME_TRY(ncnn_status(net_.load_param_mem(params_.c_str())));
  net_.load_model(reinterpret_cast<const unsigned char*>(weights_.data()));

  input_indices_ = net_.input_indexes();
  for (const auto& x : net_.input_names()) {
    input_tensors_.emplace_back(TensorDesc{
        Device("cpu"),
        DataType::kFLOAT,
        {},
        x,
    });
  }
  output_indices_ = net_.output_indexes();
  for (const auto& x : net_.output_names()) {
    output_tensors_.emplace_back(TensorDesc{
        Device("cpu"),
        DataType::kFLOAT,
        {},
        x,
    });
  }
  return success();
}

Result<void> NCNNNet::Deinit() { return success(); }

Result<void> NCNNNet::Reshape(Span<TensorShape> input_shapes) {
  for (size_t i = 0; i < input_shapes.size(); ++i) {
    input_tensors_[i].Reshape(input_shapes[i]);
  }
  return success();
}

Result<Span<Tensor>> NCNNNet::GetInputTensors() { return input_tensors_; }

Result<Span<Tensor>> NCNNNet::GetOutputTensors() { return output_tensors_; }

// TODO: discuss a policy for batch processing
Result<void> NCNNNet::Forward() {
  auto extractor = net_.create_extractor();
  OUTCOME_TRY(stream_.Wait());
  std::vector<ncnn::Mat> inputs(input_indices_.size());
  for (size_t i = 0; i < input_indices_.size(); ++i) {
    auto& tensor = input_tensors_[i];
    auto shape = tensor.shape();
    assert(shape[0] == 1);
    inputs[i] = ncnn::Mat(shape[3], shape[2], shape[1], tensor.data());
    OUTCOME_TRY(ncnn_status(extractor.input(input_indices_[i], inputs[i])));
  }
  std::vector<ncnn::Mat> outputs(output_indices_.size());
  for (size_t i = 0; i < output_indices_.size(); ++i) {
    OUTCOME_TRY(ncnn_status(extractor.extract(output_indices_[i], outputs[i])));
    auto& tensor = output_tensors_[i];
    auto shape = outputs[i].shape();
    if (outputs[i].dims == 1) {
      tensor.Reshape({1, shape.w});
    } else if (outputs[i].dims == 2) {
      tensor.Reshape({1, shape.h, shape.w});
    } else if (outputs[i].dims == 3) {
      tensor.Reshape({1, shape.d, shape.h, shape.w});
    } else {
      // for dim==4 case and blank image.
      tensor.Reshape({1, shape.c, shape.d, shape.h, shape.w});
    }
    // tensor.Reshape({1, shape.c, shape.h, shape.w});
    // ncnn Mat may be padded, flatten to avoid that
    auto flattened = outputs[i].reshape(shape.c * shape.h * shape.w);
    // if ((shape.c * shape.h * shape.w) > 0)
    if (outputs[i].dims > 0) {
      OUTCOME_TRY(tensor.CopyFrom(flattened.data, stream_));
    }
    OUTCOME_TRY(stream_.Wait());
  }
  return success();
}

class NCNNNetCreator : public Creator<Net> {
 public:
  const char* GetName() const override { return "ncnn"; }
  int GetVersion() const override { return 0; }
  std::unique_ptr<Net> Create(const Value& args) override {
    auto p = std::make_unique<NCNNNet>();
    if (auto r = p->Init(args)) {
      return p;
    } else {
      MMDEPLOY_ERROR("error creating NCNNNet: {}", r.error().message().c_str());
      return nullptr;
    }
  }
};

REGISTER_MODULE(Net, NCNNNetCreator);

}  // namespace mmdeploy