refractor module

Signed-off-by: Ceng23333 <441651826@qq.com>

Author: Ceng23333

include/infinicore/nn/module.hpp

@@ -4,11 +4,12 @@
 
 #include <unordered_map>
 #include <type_traits>
+#include <vector>
 
 namespace infinicore::nn {
 class Module {
 public:
-    Module() : path_("") {}
+    Module() = default;
 
     const std::unordered_map<std::string, Parameter> &state_dict() const;
 
@@ -18,11 +19,9 @@ class Module {
 
     void load_parameter_from_blob(const std::string &name, const void *data);
 
+protected:
     Tensor register_parameter(const std::string &name, Parameter param);
 
-    // Get the full path of this module in the hierarchy
-    const std::string &path() const { return path_; }
-
     // Add an existing submodule to this module's hierarchy
     // Template parameter M must be a type derived from Module
     // Returns the submodule for convenience (allows method chaining)
@@ -32,14 +31,6 @@ class Module {
         static_assert(std::is_base_of<Module, M>::value,
                       "Template parameter M must be derived from infinicore::nn::Module");
 
-        // Set the submodule's path based on this module's path
-        std::string submodule_path = path_.empty() ? name : path_ + "." + name;
-        submodule->path_ = submodule_path;
-
-        // Recursively update paths of all descendants
-        // This handles the case where submodules are added in any order
-        submodule->update_submodule_paths();
-
         // Store in the submodules map (std::shared_ptr<M> automatically converts to std::shared_ptr<Module>)
         submodules_[name] = submodule;
 
@@ -65,61 +56,68 @@ class Module {
     // Each submodule is named as "name.0", "name.1", etc.
     // Template parameter M must be a type derived from Module
     template <typename M, typename... Args>
-    std::vector<std::shared_ptr<M>> register_modules(size_t layers, const std::string &name, Args &&...args) {
+    std::vector<std::shared_ptr<M>> register_modules(size_t count, const std::string &name, Args &&...args) {
         static_assert(std::is_base_of<Module, M>::value,
                       "Template parameter M must be derived from infinicore::nn::Module");
 
-        auto submodules = std::vector<std::shared_ptr<M>>(layers);
-        for (size_t i = 0; i < layers; i++) {
-            submodules[i] = register_module<M>(name + "." + std::to_string(i), std::forward<Args>(args)...);
+        std::vector<std::shared_ptr<M>> modules;
+        modules.reserve(count);
+        for (size_t i = 0; i < count; i++) {
+            modules.push_back(register_module<M>(name + "." + std::to_string(i), std::forward<Args>(args)...));
         }
-        return submodules;
+        return modules;
     }
 
-private:
-    void collect_all_parameters(std::unordered_map<std::string, Parameter> &all_params) const;
-    void update_submodule_paths();
-
 protected:
-    std::string path_;  // Full path of this module in the hierarchy (e.g., "layer1.sublayer")
     Device device_;
     std::unordered_map<std::string, std::shared_ptr<Module>> submodules_;
     std::unordered_map<std::string, Parameter> parameters_;
+
+private:
+    void collect_all_parameters(std::unordered_map<std::string, Parameter> &all_params, const std::string &prefix = "") const;
 };
 
 // ============================================================================
 // PyTorch-like Macros for Convenient Module Registration
 // ============================================================================
 
 /**
- * @brief Register a submodule with automatic name inference from variable name
+ * @brief Register submodules with automatic name inference from variable name
  *
  * Usage:
  * @code
  *   class MyModel : public Module {
- *       INFINI_MODULE(Linear, layer1);
- *       INFINI_MODULE(Linear, layer2);
+ *   protected:
+ *       INFINICORE_NN_MODULE(Linear, layer1);
+ *       INFINICORE_NN_MODULE(Linear, layer2);
+ *       INFINICORE_NN_MODULE_VEC(Linear, layers);
  *
  *   public:
  *       MyModel() {
- *           INFINI_INIT_MODULE(layer1, 128, 64);
- *           INFINI_INIT_MODULE(layer2, 64, 32);
+ *           INFINICORE_NN_MODULE_INIT(layer1, 128, 64);
+ *           INFINICORE_NN_MODULE_INIT(layer2, 64, 32);
+ *           layers_ = register_modules<Linear>(3, "layers", 32, 16);
  *       }
  *   };
  * @endcode
  */
 
-// Declare a module member variable
-#define INFINI_MODULE(ModuleType, name) \
+// Declare a single module member variable
+#define INFINICORE_NN_MODULE(ModuleType, name) \
     std::shared_ptr<ModuleType> name##_
 
-// Initialize a module (use in constructor)
-#define INFINI_INIT_MODULE(name, ...) \
-    name##_ = register_module<std::remove_reference<decltype(*name##_)>::type>(#name, ##__VA_ARGS__)
+// Declare a vector of modules member variable
+#define INFINICORE_NN_MODULE_VEC(ModuleType, name) \
+    std::vector<std::shared_ptr<ModuleType>> name##_
+
+// Initialize a module in constructor
+#define INFINICORE_NN_MODULE_INIT(name, ...) \
+    name##_ = this->register_module<std::remove_reference<decltype(*name##_)>::type>(#name, ##__VA_ARGS__)
 
-// Alternative: Combined declaration and initialization helper
-// Useful when you want to initialize inline (though less flexible)
-#define INFINI_SUBMODULE(ModuleType, name, ...) \
-    name##_ = register_module<ModuleType>(#name, ##__VA_ARGS__)
+// Initialize a vector of modules in constructor
+// Usage: INFINICORE_NN_MODULE_VEC_INIT(layers, count, ModuleType, ctor_args...)
+// Example: INFINICORE_NN_MODULE_VEC_INIT(layers, 3, Linear, 128, 64)
+#define INFINICORE_NN_MODULE_VEC_INIT(name, count, ModuleType, ...) \
+    name##_ = this->register_modules<ModuleType>(count, #name, ##__VA_ARGS__)
 
 } // namespace infinicore::nn

src/infinicore/nn/module.cc

@@ -5,15 +5,15 @@ const std::unordered_map<std::string, Parameter> &Module::state_dict() const {
     static std::unordered_map<std::string, Parameter> result;
     result.clear();
 
-    collect_all_parameters(result);
+    collect_all_parameters(result, "");
 
     return result;
 }
 
 void Module::load_state_dict(const std::unordered_map<std::string, Tensor> &_state_dict) {
     // Collect all parameters from this module and its submodules with their full hierarchical names
     std::unordered_map<std::string, Parameter> all_params;
-    collect_all_parameters(all_params);
+    collect_all_parameters(all_params, "");
 
     // For each parameter in this module hierarchy, load from the state dict
     for (auto &[param_full_name, param] : all_params) {
@@ -39,28 +39,17 @@ Tensor Module::register_parameter(const std::string &name, Parameter param) {
     return param;
 }
 
-void Module::collect_all_parameters(std::unordered_map<std::string, Parameter> &all_params) const {
-    // Add direct parameters with this module's path as prefix
-    for (const auto &p : parameters_) {
-        std::string param_name = path_.empty() ? p.first : path_ + "." + p.first;
-        all_params[param_name] = p.second;
+void Module::collect_all_parameters(std::unordered_map<std::string, Parameter> &all_params, const std::string &prefix) const {
+    // Add direct parameters with the given prefix
+    for (const auto &[param_name, param] : parameters_) {
+        std::string full_name = prefix.empty() ? param_name : prefix + "." + param_name;
+        all_params[full_name] = param;
     }
 
-    // Recursively collect parameters from submodules
-    // Each submodule already knows its own path, so no need to pass prefix
-    for (const auto &sub : submodules_) {
-        sub.second->collect_all_parameters(all_params);
-    }
-}
-
-void Module::update_submodule_paths() {
-    // Recursively update the paths of all submodules based on this module's current path
-    for (auto &[name, submodule] : submodules_) {
-        // Update submodule's path based on this module's current path
-        std::string submodule_path = path_.empty() ? name : path_ + "." + name;
-        submodule->path_ = submodule_path;
-        // Recursively update descendants
-        submodule->update_submodule_paths();
+    // Recursively collect parameters from submodules with extended prefix
+    for (const auto &[sub_name, submodule] : submodules_) {
+        std::string sub_prefix = prefix.empty() ? sub_name : prefix + "." + sub_name;
+        submodule->collect_all_parameters(all_params, sub_prefix);
     }
 }