Add NASNet-a model

bin/export-model-arch/src/export-model-arch/main.cc

@@ -5,6 +5,7 @@
 #include "models/candle_uno/candle_uno.h"
 #include "models/dlrm/dlrm.h"
 #include "models/inception_v3/inception_v3.h"
+#include "models/nasnet/nasnet.h"
 #include "models/split_test/split_test.h"
 #include "models/transformer/transformer.h"
 #include "models/yolov10/yolov10.h"
@@ -67,6 +68,9 @@ tl::expected<ComputationGraph, std::string>
   } else if (model_name == "inception_v3") {
     return get_inception_v3_computation_graph(
         get_default_inception_v3_training_config());
+  } else if (model_name == "nasnet_a_large") {
+    return get_nasnet_a_large_computation_graph(
+        get_default_nasnet_a_large_config());
   } else if (model_name == "candle_uno") {
     return get_candle_uno_computation_graph(get_default_candle_uno_config());
   } else if (model_name == "bert") {
@@ -154,6 +158,7 @@ int main(int argc, char **argv) {
   std::vector<std::string> model_options = {
       "transformer",
       "inception_v3",
+      "nasnet_a_large",
       "candle_uno",
       "bert",
       "dlrm",

lib/models/include/models/nasnet/index.dox

@@ -0,0 +1,11 @@
+namespace FlexFlow {
+/**
+
+\page models-nasnet models/nasnet
+
+NASNet-A Large image classification model based on the implementation in
+<a href="https://github.com/huggingface/pytorch-image-models/blob/7cb8dca52b47a6bab108434c3b088d36a11feb1a/timm/models/nasnet.py">
+timm</a>.
+
+*/
+}

lib/models/include/models/nasnet/nasnet.h

@@ -0,0 +1,31 @@
+/**
+ * @file nasnet.h
+ *
+ * @brief NASNet model computation graphs.
+ */
+
+#ifndef _FLEXFLOW_LIB_MODELS_INCLUDE_MODELS_NASNET_NASNET_H
+#define _FLEXFLOW_LIB_MODELS_INCLUDE_MODELS_NASNET_NASNET_H
+
+#include "models/nasnet/nasnet_a_large_config.dtg.h"
+#include "pcg/computation_graph.dtg.h"
+
+namespace FlexFlow {
+
+/**
+ * @brief Get the default NASNet-A Large configuration.
+ *
+ * The defaults match the timm NASNet-A Large definition: RGB 331x331 input,
+ * 1000 output classes, 4032 final features, and no classifier dropout.
+ */
+NASNetALargeConfig get_default_nasnet_a_large_config();
+
+/**
+ * @brief Build a NASNet-A Large computation graph.
+ */
+ComputationGraph
+    get_nasnet_a_large_computation_graph(NASNetALargeConfig const &config);
+
+} // namespace FlexFlow
+
+#endif

lib/models/include/models/nasnet/nasnet_a_large_config.dtg.toml

@@ -0,0 +1,44 @@
+namespace = "FlexFlow"
+name = "NASNetALargeConfig"
+type = "struct"
+
+features = [
+  "eq",
+  "ord",
+  "hash",
+  "json",
+  "rapidcheck",
+  "fmt",
+]
+
+includes = [
+  "utils/positive_int/positive_int.h",
+]
+
+docstring = '''
+\brief Configuration for a NASNet-A Large image classifier.
+'''
+
+[[fields]]
+name = "batch_size"
+type = "::FlexFlow::positive_int"
+
+[[fields]]
+name = "num_classes"
+type = "::FlexFlow::positive_int"
+
+[[fields]]
+name = "input_channels"
+type = "::FlexFlow::positive_int"
+
+[[fields]]
+name = "image_height"
+type = "::FlexFlow::positive_int"
+
+[[fields]]
+name = "image_width"
+type = "::FlexFlow::positive_int"
+
+[[fields]]
+name = "dropout"
+type = "float"

lib/models/index.dox

@@ -10,6 +10,7 @@ namespace FlexFlow {
 - \subpage models-candle-uno "Candle UNO"
 - \subpage models-dlrm "DLRM"
 - \subpage models-inception "Inception v3"
+- \subpage models-nasnet "NASNet-A Large"
 - \subpage models-transformer "Transformer"
 
 \section test-models Artificial Models for Testing