C++ Models (#728)

* Added the existing code

* Added squeezenet and fixed some stuff in the other models

* Wrote DenseNet and a part of InceptionV3

Going to clean and check all of the models and finish inception

* Fixed some errors in the models

Next step is writing inception and comparing with python code again.

* Completed inception and changed models directory

* Fixed and wrote some stuff

* fixed maxpoool2d and avgpool2d and adaptiveavgpool2d

* Fixed a few stuff

Moved cmakelists to root and changed the namespace to vision and wrote weight initialization in inception

* Added models namespace and changed cmakelists

the project is now installable

* Removed some comments

* Changed style to pytorch style, added some comments and fixed some minor errors

* Removed truncated normal init

* Changed classes to structs and fixed a few errors

* Replaced modelsimpl structs with functional wherever possible

* Changed adaptive average pool from struct to function

* Wrote a max_pool2d wrapper and added some comments

* Replaced xavier init with kaiming init

* Fixed an error in kaiming inits

* Added model conversion and tests

* Fixed a typo in alexnet and removed tests from cmake

* Made an extension of tests and added module names to Densenet

* Added python tests

* Added MobileNet and GoogLeNet models

* Added tests and conversions for new models and fixed a few errors

* Updated Alexnet ad VGG

* Updated Densenet, Squeezenet and Inception

* Added ResNexts and their conversions

* Added tests for ResNexts

* Wrote tools nessesary to write ShuffleNet

* Added ShuffleNetV2

* Fixed some errors in ShuffleNetV2

* Added conversions for shufflenetv2

* Fixed the errors in test_models.cpp

* Updated setup.py

* Fixed flake8 error on test_cpp_models.py

* Changed view to reshape in forward of ResNet

* Updated ShuffleNetV2

* Split extensions to tests and ops

* Fixed test extension

* Fixed image path in test_cpp_models.py

* Fixed image path in test_cpp_models.py

* Fixed a few things in test_cpp_models.py

* Put the test models in evaluation mode

* Fixed registering error in GoogLeNet

* Updated setup.py

* write test_cpp_models.py with unittest

* Fixed a problem with pytest in test_cpp_models.py

* Fixed a lint problem

torchvision/csrc/models/shufflenetv2.h

+#ifndef SHUFFLENETV2_H
+#define SHUFFLENETV2_H
+
+#include <torch/torch.h>
+
+namespace vision {
+namespace models {
+
+struct ShuffleNetV2Impl : torch::nn::Module {
+  std::vector<int64_t> _stage_out_channels;
+  torch::nn::Sequential conv1{nullptr}, stage2, stage3, stage4, conv5{nullptr};
+  torch::nn::Linear fc{nullptr};
+
+  ShuffleNetV2Impl(
+      const std::vector<int64_t>& stage_repeats,
+      const std::vector<int64_t>& stage_out_channels,
+      int64_t num_classes = 1000);
+
+  torch::Tensor forward(torch::Tensor x);
+};
+
+struct ShuffleNetV2_x0_5Impl : ShuffleNetV2Impl {
+  ShuffleNetV2_x0_5Impl(int64_t num_classes = 1000);
+};
+
+struct ShuffleNetV2_x1_0Impl : ShuffleNetV2Impl {
+  ShuffleNetV2_x1_0Impl(int64_t num_classes = 1000);
+};
+
+struct ShuffleNetV2_x1_5Impl : ShuffleNetV2Impl {
+  ShuffleNetV2_x1_5Impl(int64_t num_classes = 1000);
+};
+
+struct ShuffleNetV2_x2_0Impl : ShuffleNetV2Impl {
+  ShuffleNetV2_x2_0Impl(int64_t num_classes = 1000);
+};
+
+TORCH_MODULE(ShuffleNetV2);
+TORCH_MODULE(ShuffleNetV2_x0_5);
+TORCH_MODULE(ShuffleNetV2_x1_0);
+TORCH_MODULE(ShuffleNetV2_x1_5);
+TORCH_MODULE(ShuffleNetV2_x2_0);
+
+} // namespace models
+} // namespace vision
+
+#endif // SHUFFLENETV2_H

torchvision/csrc/models/squeezenet.cpp

+#include "squeezenet.h"
+
+#include <limits>
+#include "modelsimpl.h"
+
+namespace vision {
+namespace models {
+struct Fire : torch::nn::Module {
+  torch::nn::Conv2d squeeze, expand1x1, expand3x3;
+
+  Fire(
+      int64_t inplanes,
+      int64_t squeeze_planes,
+      int64_t expand1x1_planes,
+      int64_t expand3x3_planes)
+      : squeeze(torch::nn::Conv2dOptions(inplanes, squeeze_planes, 1)),
+        expand1x1(
+            torch::nn::Conv2dOptions(squeeze_planes, expand1x1_planes, 1)),
+        expand3x3(torch::nn::Conv2dOptions(squeeze_planes, expand3x3_planes, 3)
+                      .padding(1)) {
+    register_module("squeeze", squeeze);
+    register_module("expand1x1", expand1x1);
+    register_module("expand3x3", expand3x3);
+  }
+
+  torch::Tensor forward(torch::Tensor x) {
+    x = torch::relu(squeeze->forward(x));
+    return torch::cat(
+        {torch::relu(expand1x1->forward(x)),
+         torch::relu(expand3x3->forward(x))},
+        1);
+  }
+};
+
+SqueezeNetImpl::SqueezeNetImpl(double version, int64_t num_classes)
+    : num_classes(num_classes) {
+  if (modelsimpl::double_compare(version, 1.0)) {
+    features = torch::nn::Sequential(
+        torch::nn::Conv2d(torch::nn::Conv2dOptions(3, 96, 7).stride(2)),
+        torch::nn::Functional(modelsimpl::relu_),
+        torch::nn::Functional(torch::max_pool2d, 3, 2, 0, 1, true),
+        Fire(96, 16, 64, 64),
+        Fire(128, 16, 64, 64),
+        Fire(128, 32, 128, 128),
+        torch::nn::Functional(torch::max_pool2d, 3, 2, 0, 1, true),
+        Fire(256, 32, 128, 128),
+        Fire(256, 48, 192, 192),
+        Fire(384, 48, 192, 192),
+        Fire(384, 64, 256, 256),
+        torch::nn::Functional(torch::max_pool2d, 3, 2, 0, 1, true),
+        Fire(512, 64, 256, 256));
+  } else if (modelsimpl::double_compare(version, 1.1)) {
+    features = torch::nn::Sequential(
+        torch::nn::Conv2d(torch::nn::Conv2dOptions(3, 64, 3).stride(2)),
+        torch::nn::Functional(modelsimpl::relu_),
+        torch::nn::Functional(torch::max_pool2d, 3, 2, 0, 1, true),
+        Fire(64, 16, 64, 64),
+        Fire(128, 16, 64, 64),
+        torch::nn::Functional(torch::max_pool2d, 3, 2, 0, 1, true),
+        Fire(128, 32, 128, 128),
+        Fire(256, 32, 128, 128),
+        torch::nn::Functional(torch::max_pool2d, 3, 2, 0, 1, true),
+        Fire(256, 48, 192, 192),
+        Fire(384, 48, 192, 192),
+        Fire(384, 64, 256, 256),
+        Fire(512, 64, 256, 256));
+  } else {
+    std::cerr << "Wrong version number is passed th SqueeseNet constructor!"
+              << std::endl;
+    assert(false);
+  }
+
+  // Final convolution is initialized differently from the rest
+  auto final_conv =
+      torch::nn::Conv2d(torch::nn::Conv2dOptions(512, num_classes, 1));
+
+  classifier = torch::nn::Sequential(
+      torch::nn::Dropout(0.5),
+      final_conv,
+      torch::nn::Functional(modelsimpl::relu_),
+      torch::nn::Functional(modelsimpl::adaptive_avg_pool2d, 1));
+
+  register_module("features", features);
+  register_module("classifier", classifier);
+
+  for (auto& module : modules(/*include_self=*/false))
+    if (auto M = dynamic_cast<torch::nn::Conv2dImpl*>(module.get())) {
+      if (M == final_conv.get())
+        torch::nn::init::normal_(M->weight, 0.0, 0.01);
+      else
+        torch::nn::init::kaiming_uniform_(M->weight);
+
+      if (M->options.with_bias())
+        torch::nn::init::constant_(M->bias, 0);
+    }
+}
+
+torch::Tensor SqueezeNetImpl::forward(torch::Tensor x) {
+  x = features->forward(x);
+  x = classifier->forward(x);
+  return x.view({x.size(0), -1});
+}
+
+SqueezeNet1_0Impl::SqueezeNet1_0Impl(int64_t num_classes)
+    : SqueezeNetImpl(1.0, num_classes) {}
+
+SqueezeNet1_1Impl::SqueezeNet1_1Impl(int64_t num_classes)
+    : SqueezeNetImpl(1.1, num_classes) {}
+
+} // namespace models
+} // namespace vision

torchvision/csrc/models/squeezenet.h

+#ifndef SQUEEZENET_H
+#define SQUEEZENET_H
+
+#include <torch/torch.h>
+
+namespace vision {
+namespace models {
+struct SqueezeNetImpl : torch::nn::Module {
+  int64_t num_classes;
+  torch::nn::Sequential features{nullptr}, classifier{nullptr};
+
+  SqueezeNetImpl(double version = 1.0, int64_t num_classes = 1000);
+
+  torch::Tensor forward(torch::Tensor x);
+};
+
+// SqueezeNet model architecture from the "SqueezeNet: AlexNet-level
+// accuracy with 50x fewer parameters and <0.5MB model size"
+// <https://arxiv.org/abs/1602.07360> paper.
+struct SqueezeNet1_0Impl : SqueezeNetImpl {
+  SqueezeNet1_0Impl(int64_t num_classes = 1000);
+};
+
+// SqueezeNet 1.1 model from the official SqueezeNet repo
+// <https://github.com/DeepScale/SqueezeNet/tree/master/SqueezeNet_v1.1>.
+// SqueezeNet 1.1 has 2.4x less computation and slightly fewer parameters
+// than SqueezeNet 1.0, without sacrificing accuracy.
+struct SqueezeNet1_1Impl : SqueezeNetImpl {
+  SqueezeNet1_1Impl(int64_t num_classes = 1000);
+};
+
+TORCH_MODULE(SqueezeNet);
+TORCH_MODULE(SqueezeNet1_0);
+TORCH_MODULE(SqueezeNet1_1);
+
+} // namespace models
+} // namespace vision
+
+#endif // SQUEEZENET_H

torchvision/csrc/models/vgg.cpp

+#include "vgg.h"
+
+#include <unordered_map>
+#include "modelsimpl.h"
+
+namespace vision {
+namespace models {
+torch::nn::Sequential makeLayers(
+    const std::vector<int>& cfg,
+    bool batch_norm = false) {
+  torch::nn::Sequential seq;
+  auto channels = 3;
+
+  for (const auto& V : cfg) {
+    if (V <= -1)
+      seq->push_back(torch::nn::Functional(modelsimpl::max_pool2d, 2, 2));
+    else {
+      seq->push_back(torch::nn::Conv2d(
+          torch::nn::Conv2dOptions(channels, V, 3).padding(1)));
+
+      if (batch_norm)
+        seq->push_back(torch::nn::BatchNorm(V));
+      seq->push_back(torch::nn::Functional(modelsimpl::relu_));
+
+      channels = V;
+    }
+  }
+
+  return seq;
+}
+
+void VGGImpl::_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,
+          /*a=*/0,
+          torch::nn::init::FanMode::FanOut,
+          torch::nn::init::Nonlinearity::ReLU);
+      torch::nn::init::constant_(M->bias, 0);
+    } else if (auto M = dynamic_cast<torch::nn::BatchNormImpl*>(module.get())) {
+      torch::nn::init::constant_(M->weight, 1);
+      torch::nn::init::constant_(M->bias, 0);
+    } else if (auto M = dynamic_cast<torch::nn::LinearImpl*>(module.get())) {
+      torch::nn::init::normal_(M->weight, 0, 0.01);
+      torch::nn::init::constant_(M->bias, 0);
+    }
+  }
+}
+
+VGGImpl::VGGImpl(
+    torch::nn::Sequential features,
+    int64_t num_classes,
+    bool initialize_weights) {
+  classifier = torch::nn::Sequential(
+      torch::nn::Linear(512 * 7 * 7, 4096),
+      torch::nn::Functional(modelsimpl::relu_),
+      torch::nn::Dropout(),
+      torch::nn::Linear(4096, 4096),
+      torch::nn::Functional(modelsimpl::relu_),
+      torch::nn::Dropout(),
+      torch::nn::Linear(4096, num_classes));
+
+  this->features = features;
+
+  register_module("features", this->features);
+  register_module("classifier", classifier);
+
+  if (initialize_weights)
+    _initialize_weights();
+}
+
+torch::Tensor VGGImpl::forward(torch::Tensor x) {
+  x = features->forward(x);
+  x = torch::adaptive_avg_pool2d(x, {7, 7});
+  x = x.view({x.size(0), -1});
+  x = classifier->forward(x);
+  return x;
+}
+
+// clang-format off
+static std::unordered_map<char, std::vector<int>> cfg = {
+  {'A', {64, -1, 128, -1, 256, 256, -1, 512, 512, -1, 512, 512, -1}},
+  {'B', {64, 64, -1, 128, 128, -1, 256, 256, -1, 512, 512, -1, 512, 512, -1}},
+  {'D', {64, 64, -1, 128, 128, -1, 256, 256, 256, -1, 512, 512, 512, -1, 512, 512, 512, -1}},
+  {'E', {64,  64,  -1,  128, 128, -1,  256, 256, 256, 256, -1, 512, 512, 512, 512, -1,  512, 512, 512, 512, -1}}};
+// clang-format on
+
+VGG11Impl::VGG11Impl(int64_t num_classes, bool initialize_weights)
+    : VGGImpl(makeLayers(cfg['A']), num_classes, initialize_weights) {}
+
+VGG13Impl::VGG13Impl(int64_t num_classes, bool initialize_weights)
+    : VGGImpl(makeLayers(cfg['B']), num_classes, initialize_weights) {}
+
+VGG16Impl::VGG16Impl(int64_t num_classes, bool initialize_weights)
+    : VGGImpl(makeLayers(cfg['D']), num_classes, initialize_weights) {}
+
+VGG19Impl::VGG19Impl(int64_t num_classes, bool initialize_weights)
+    : VGGImpl(makeLayers(cfg['E']), num_classes, initialize_weights) {}
+
+VGG11BNImpl::VGG11BNImpl(int64_t num_classes, bool initialize_weights)
+    : VGGImpl(makeLayers(cfg['A'], true), num_classes, initialize_weights) {}
+
+VGG13BNImpl::VGG13BNImpl(int64_t num_classes, bool initialize_weights)
+    : VGGImpl(makeLayers(cfg['B'], true), num_classes, initialize_weights) {}
+
+VGG16BNImpl::VGG16BNImpl(int64_t num_classes, bool initialize_weights)
+    : VGGImpl(makeLayers(cfg['D'], true), num_classes, initialize_weights) {}
+
+VGG19BNImpl::VGG19BNImpl(int64_t num_classes, bool initialize_weights)
+    : VGGImpl(makeLayers(cfg['E'], true), num_classes, initialize_weights) {}
+
+} // namespace models
+} // namespace vision

torchvision/csrc/models/vgg.h

+#ifndef VGG_H
+#define VGG_H
+
+#include <torch/torch.h>
+
+namespace vision {
+namespace models {
+struct VGGImpl : torch::nn::Module {
+  torch::nn::Sequential features{nullptr}, classifier{nullptr};
+
+  void _initialize_weights();
+
+  VGGImpl(
+      torch::nn::Sequential features,
+      int64_t num_classes = 1000,
+      bool initialize_weights = true);
+
+  torch::Tensor forward(torch::Tensor x);
+};
+
+// VGG 11-layer model (configuration "A")
+struct VGG11Impl : VGGImpl {
+  VGG11Impl(int64_t num_classes = 1000, bool initialize_weights = true);
+};
+
+// VGG 13-layer model (configuration "B")
+struct VGG13Impl : VGGImpl {
+  VGG13Impl(int64_t num_classes = 1000, bool initialize_weights = true);
+};
+
+// VGG 16-layer model (configuration "D")
+struct VGG16Impl : VGGImpl {
+  VGG16Impl(int64_t num_classes = 1000, bool initialize_weights = true);
+};
+
+// VGG 19-layer model (configuration "E")
+struct VGG19Impl : VGGImpl {
+  VGG19Impl(int64_t num_classes = 1000, bool initialize_weights = true);
+};
+
+// VGG 11-layer model (configuration "A") with batch normalization
+struct VGG11BNImpl : VGGImpl {
+  VGG11BNImpl(int64_t num_classes = 1000, bool initialize_weights = true);
+};
+
+// VGG 13-layer model (configuration "B") with batch normalization
+struct VGG13BNImpl : VGGImpl {
+  VGG13BNImpl(int64_t num_classes = 1000, bool initialize_weights = true);
+};
+
+// VGG 16-layer model (configuration "D") with batch normalization
+struct VGG16BNImpl : VGGImpl {
+  VGG16BNImpl(int64_t num_classes = 1000, bool initialize_weights = true);
+};
+
+// VGG 19-layer model (configuration 'E') with batch normalization
+struct VGG19BNImpl : VGGImpl {
+  VGG19BNImpl(int64_t num_classes = 1000, bool initialize_weights = true);
+};
+
+TORCH_MODULE(VGG);
+
+TORCH_MODULE(VGG11);
+TORCH_MODULE(VGG13);
+TORCH_MODULE(VGG16);
+TORCH_MODULE(VGG19);
+
+TORCH_MODULE(VGG11BN);
+TORCH_MODULE(VGG13BN);
+TORCH_MODULE(VGG16BN);
+TORCH_MODULE(VGG19BN);
+
+} // namespace models
+} // namespace vision
+
+#endif // VGG_H

torchvision/csrc/vision.h

+#ifndef VISION_H
+#define VISION_H
+
+#include <torchvision/models/models.h>
+
+#endif  // VISION_H