Mnasnet Model for C++ Side (#1108)

* Added mnasnet

* Fixed some stuff

* Fixed some stuff

* Finished MNASNet

* Fixed format error

* Fixed format error

test/test_cpp_models.py

@@ -113,6 +113,18 @@ class Tester(unittest.TestCase):
     def test_googlenet(self):
         process_model(models.googlenet(self.pretrained), self.image, _C_tests.forward_googlenet, 'GoogLeNet')
 
+    def test_mnasnet0_5(self):
+        process_model(models.mnasnet0_5(self.pretrained), self.image, _C_tests.forward_mnasnet0_5, 'MNASNet0_5')
+
+    def test_mnasnet0_75(self):
+        process_model(models.mnasnet0_75(self.pretrained), self.image, _C_tests.forward_mnasnet0_75, 'MNASNet0_75')
+
+    def test_mnasnet1_0(self):
+        process_model(models.mnasnet1_0(self.pretrained), self.image, _C_tests.forward_mnasnet1_0, 'MNASNet1_0')
+
+    def test_mnasnet1_3(self):
+        process_model(models.mnasnet1_3(self.pretrained), self.image, _C_tests.forward_mnasnet1_3, 'MNASNet1_3')
+
     def test_inception_v3(self):
         self.image = read_image2()
         process_model(models.inception_v3(self.pretrained), self.image, _C_tests.forward_inceptionv3, 'Inceptionv3')

test/test_models.cpp

@@ -129,6 +129,19 @@ torch::Tensor forward_inceptionv3(
   return network->forward(x).output;
 }
 
+torch::Tensor forward_mnasnet0_5(const std::string& input_path, torch::Tensor x) {
+  return forward_model<MNASNet0_5>(input_path, x);
+}
+torch::Tensor forward_mnasnet0_75(const std::string& input_path, torch::Tensor x) {
+  return forward_model<MNASNet0_75>(input_path, x);
+}
+torch::Tensor forward_mnasnet1_0(const std::string& input_path, torch::Tensor x) {
+  return forward_model<MNASNet1_0>(input_path, x);
+}
+torch::Tensor forward_mnasnet1_3(const std::string& input_path, torch::Tensor x) {
+  return forward_model<MNASNet1_3>(input_path, x);
+}
+
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   m.def("forward_alexnet", &forward_alexnet, "forward_alexnet");
 
@@ -170,4 +183,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
 
   m.def("forward_googlenet", &forward_googlenet, "forward_googlenet");
   m.def("forward_inceptionv3", &forward_inceptionv3, "forward_inceptionv3");
+
+  m.def("forward_mnasnet0_5", &forward_mnasnet0_5, "forward_mnasnet0_5");
+  m.def("forward_mnasnet0_75", &forward_mnasnet0_75, "forward_mnasnet0_75");
+  m.def("forward_mnasnet1_0", &forward_mnasnet1_0, "forward_mnasnet1_0");
+  m.def("forward_mnasnet1_3", &forward_mnasnet1_3, "forward_mnasnet1_3");
 }

torchvision/csrc/convert_models/convert_models.cpp

@@ -72,5 +72,14 @@ int main(int argc, const char* argv[]) {
   convert_and_save_model<InceptionV3>(
       "inceptionv3_python.pt", "inceptionv3_cpp.pt");
 
+  convert_and_save_model<MNASNet0_5>(
+      "mnasnet0_5_python.pt", "mnasnet0_5_cpp.pt");
+  convert_and_save_model<MNASNet0_75>(
+      "mnasnet0_75_python.pt", "mnasnet0_75_cpp.pt");
+  convert_and_save_model<MNASNet1_0>(
+      "mnasnet1_0_python.pt", "mnasnet1_0_cpp.pt");
+  convert_and_save_model<MNASNet1_3>(
+      "mnasnet1_3_python.pt", "mnasnet1_3_cpp.pt");
+
   return 0;
 }

torchvision/csrc/models/mnasnet.cpp

+#include "mnasnet.h"
+
+#include "modelsimpl.h"
+
+namespace vision {
+namespace models {
+using Options = torch::nn::Conv2dOptions;
+
+struct MNASNetInvertedResidualImpl : torch::nn::Module {
+  bool apply_residual;
+  torch::nn::Sequential layers;
+
+  MNASNetInvertedResidualImpl(
+      int64_t input,
+      int64_t output,
+      int64_t kernel,
+      int64_t stride,
+      double expansion_factor,
+      double bn_momentum = 0.1) {
+    assert(stride == 1 || stride == 2);
+    assert(kernel == 3 || kernel == 5);
+
+    auto mid = int64_t(input * expansion_factor);
+    apply_residual = input == output && stride == 1;
+
+    layers->push_back(
+        torch::nn::Conv2d(Options(input, mid, 1).with_bias(false)));
+    layers->push_back(torch::nn::BatchNorm(
+        torch::nn::BatchNormOptions(mid).momentum(bn_momentum)));
+    layers->push_back(
+        torch::nn::Functional(torch::nn::Functional(modelsimpl::relu_)));
+    layers->push_back(
+        torch::nn::Conv2d(torch::nn::Conv2d(Options(mid, mid, kernel)
+                                                .padding(kernel / 2)
+                                                .stride(stride)
+                                                .groups(mid)
+                                                .with_bias(false))));
+    layers->push_back(torch::nn::BatchNorm(
+        torch::nn::BatchNormOptions(mid).momentum(bn_momentum)));
+    layers->push_back(
+        torch::nn::Functional(torch::nn::Functional(modelsimpl::relu_)));
+    layers->push_back(
+        torch::nn::Conv2d(Options(mid, output, 1).with_bias(false)));
+    layers->push_back(torch::nn::BatchNorm(
+        torch::nn::BatchNormOptions(output).momentum(bn_momentum)));
+
+    register_module("layers", layers);
+  }
+
+  torch::Tensor forward(torch::Tensor x) {
+    if (apply_residual)
+      return layers->forward(x) + x;
+    return layers->forward(x);
+  }
+};
+
+TORCH_MODULE(MNASNetInvertedResidual);
+
+struct StackSequentailImpl : torch::nn::SequentialImpl {
+  using SequentialImpl::SequentialImpl;
+
+  torch::Tensor forward(torch::Tensor x) {
+    return SequentialImpl::forward(x);
+  }
+};
+
+TORCH_MODULE(StackSequentail);
+
+StackSequentail stack(
+    int64_t input,
+    int64_t output,
+    int64_t kernel,
+    int64_t stride,
+    double exp_factor,
+    int64_t repeats,
+    double bn_momentum) {
+  assert(repeats >= 1);
+
+  StackSequentail seq;
+  seq->push_back(MNASNetInvertedResidual(
+      input, output, kernel, stride, exp_factor, bn_momentum));
+
+  for (int64_t i = 1; i < repeats; ++i)
+    seq->push_back(MNASNetInvertedResidual(
+        output, output, kernel, 1, exp_factor, bn_momentum));
+
+  return seq;
+}
+
+int64_t round_to_multiple_of(
+    int64_t val,
+    int64_t divisor,
+    double round_up_bias = .9) {
+  assert(0.0 < round_up_bias && round_up_bias < 1.0);
+  auto new_val = std::max(divisor, (val + divisor / 2) / divisor * divisor);
+  return new_val >= round_up_bias * val ? new_val : new_val + divisor;
+}
+
+std::vector<int64_t> scale_depths(std::vector<int64_t> depths, double alpha) {
+  std::vector<int64_t> data(depths.size());
+  for (size_t i = 0; i < data.size(); ++i) {
+    data[i] = round_to_multiple_of(int64_t(depths[i] * alpha), 8);
+  }
+
+  return data;
+}
+
+void MNASNetImpl::_initialize_weights() {
+  for (auto& module : modules(/*include_self=*/false)) {
+    if (auto M = dynamic_cast<torch::nn::Conv2dImpl*>(module.get()))
+      torch::nn::init::kaiming_normal_(
+          M->weight,
+          0,
+          torch::nn::init::FanMode::FanOut,
+          torch::nn::init::Nonlinearity::ReLU);
+    else if (auto M = dynamic_cast<torch::nn::BatchNormImpl*>(module.get())) {
+      torch::nn::init::ones_(M->weight);
+      torch::nn::init::zeros_(M->bias);
+    } else if (auto M = dynamic_cast<torch::nn::LinearImpl*>(module.get())) {
+      torch::nn::init::normal_(M->weight, 0, 0.01);
+      torch::nn::init::zeros_(M->bias);
+    }
+  }
+}
+
+#define BN_MOMENTUM 1 - 0.9997
+
+MNASNetImpl::MNASNetImpl(double alpha, int64_t num_classes, double dropout) {
+  auto depths = scale_depths({24, 40, 80, 96, 192, 320}, alpha);
+
+  layers->push_back(torch::nn::Conv2d(
+      Options(3, 32, 3).padding(1).stride(2).with_bias(false)));
+  layers->push_back(torch::nn::BatchNorm(
+      torch::nn::BatchNormOptions(32).momentum(BN_MOMENTUM)));
+  layers->push_back(torch::nn::Functional(modelsimpl::relu_));
+  layers->push_back(torch::nn::Conv2d(
+      Options(32, 32, 3).padding(1).stride(1).groups(32).with_bias(false)));
+  layers->push_back(torch::nn::BatchNorm(
+      torch::nn::BatchNormOptions(32).momentum(BN_MOMENTUM)));
+  layers->push_back(torch::nn::Functional(modelsimpl::relu_));
+  layers->push_back(torch::nn::Conv2d(
+      Options(32, 16, 1).padding(0).stride(1).with_bias(false)));
+  layers->push_back(torch::nn::BatchNorm(
+      torch::nn::BatchNormOptions(16).momentum(BN_MOMENTUM)));
+
+  layers->push_back(stack(16, depths[0], 3, 2, 3, 3, BN_MOMENTUM));
+  layers->push_back(stack(depths[0], depths[1], 5, 2, 3, 3, BN_MOMENTUM));
+  layers->push_back(stack(depths[1], depths[2], 5, 2, 6, 3, BN_MOMENTUM));
+  layers->push_back(stack(depths[2], depths[3], 3, 1, 6, 2, BN_MOMENTUM));
+  layers->push_back(stack(depths[3], depths[4], 5, 2, 6, 4, BN_MOMENTUM));
+  layers->push_back(stack(depths[4], depths[5], 3, 1, 6, 1, BN_MOMENTUM));
+
+  layers->push_back(torch::nn::Conv2d(
+      Options(depths[5], 1280, 1).padding(0).stride(1).with_bias(false)));
+  layers->push_back(torch::nn::BatchNorm(
+      torch::nn::BatchNormOptions(1280).momentum(BN_MOMENTUM)));
+  layers->push_back(torch::nn::Functional(modelsimpl::relu_));
+
+  classifier = torch::nn::Sequential(
+      torch::nn::Dropout(dropout), torch::nn::Linear(1280, num_classes));
+
+  register_module("layers", layers);
+  register_module("classifier", classifier);
+
+  _initialize_weights();
+}
+
+torch::Tensor MNASNetImpl::forward(torch::Tensor x) {
+  x = layers->forward(x);
+  x = x.mean({2, 3});
+  return classifier->forward(x);
+}
+
+MNASNet0_5Impl::MNASNet0_5Impl(int64_t num_classes, double dropout)
+    : MNASNetImpl(.5, num_classes, dropout) {}
+
+MNASNet0_75Impl::MNASNet0_75Impl(int64_t num_classes, double dropout)
+    : MNASNetImpl(.75, num_classes, dropout) {}
+
+MNASNet1_0Impl::MNASNet1_0Impl(int64_t num_classes, double dropout)
+    : MNASNetImpl(1, num_classes, dropout) {}
+
+MNASNet1_3Impl::MNASNet1_3Impl(int64_t num_classes, double dropout)
+    : MNASNetImpl(1.3, num_classes, dropout) {}
+
+} // namespace models
+} // namespace vision

torchvision/csrc/models/mnasnet.h

+#ifndef MNASNET_H
+#define MNASNET_H
+
+#include <torch/torch.h>
+#include "general.h"
+
+namespace vision {
+namespace models {
+struct VISION_API MNASNetImpl : torch::nn::Module {
+  torch::nn::Sequential layers, classifier;
+
+  void _initialize_weights();
+
+  MNASNetImpl(double alpha, int64_t num_classes = 1000, double dropout = .2);
+
+  torch::Tensor forward(torch::Tensor x);
+};
+
+struct MNASNet0_5Impl : MNASNetImpl {
+  MNASNet0_5Impl(int64_t num_classes = 1000, double dropout = .2);
+};
+
+struct MNASNet0_75Impl : MNASNetImpl {
+  MNASNet0_75Impl(int64_t num_classes = 1000, double dropout = .2);
+};
+
+struct MNASNet1_0Impl : MNASNetImpl {
+  MNASNet1_0Impl(int64_t num_classes = 1000, double dropout = .2);
+};
+
+struct MNASNet1_3Impl : MNASNetImpl {
+  MNASNet1_3Impl(int64_t num_classes = 1000, double dropout = .2);
+};
+
+TORCH_MODULE(MNASNet);
+TORCH_MODULE(MNASNet0_5);
+TORCH_MODULE(MNASNet0_75);
+TORCH_MODULE(MNASNet1_0);
+TORCH_MODULE(MNASNet1_3);
+
+} // namespace models
+} // namespace vision
+
+#endif // MNASNET_H

torchvision/csrc/models/models.h

@@ -5,6 +5,7 @@
 #include "densenet.h"
 #include "googlenet.h"
 #include "inception.h"
+#include "mnasnet.h"
 #include "mobilenet.h"
 #include "resnet.h"
 #include "shufflenetv2.h"