issue/517 c++ nn::module

include/infinicore/nn/module.hpp

+#pragma once
+
+#include "parameter.hpp"
+
+#include <unordered_map>
+
+namespace infinicore::nn {
+class Module {
+public:
+    const std::unordered_map<std::string, Parameter> &state_dict() const;
+
+    void load_state_dict(const std::unordered_map<std::string, Tensor> &_state_dict);
+
+    void load_parameter(const std::string &name, const Tensor &param);
+
+    void load_parameter_from_blob(const std::string &name, const void *data);
+
+    Tensor register_parameter(const std::string &name, Parameter param);
+
+    template <typename M>
+    std::shared_ptr<M> add_module(const std::string &name, std::shared_ptr<M> submodule) {
+        submodules_[name] = submodule;
+        for (auto &p : submodule->parameters_) {
+            parameters_[name + "." + p.first] = p.second;
+        }
+        return submodule;
+    }
+
+    template <typename M, typename... Args>
+    std::shared_ptr<M> register_module(const std::string &name, Args &&...args) {
+        auto submodule = std::make_shared<M>(std::forward<Args>(args)...);
+        return add_module(name, submodule);
+    }
+
+    template <typename M, typename... Args>
+    std::vector<std::shared_ptr<M>> register_modules(size_t layers, const std::string &name, Args &&...args) {
+        auto submodules = std::vector<std::shared_ptr<M>>(layers);
+        for (size_t i = 0; i < layers; i++) {
+            register_module<M>(name + "." + std::to_string(i), std::forward<Args>(args)...);
+        }
+        return submodules;
+    }
+
+protected:
+    Device device_;
+    std::unordered_map<std::string, std::shared_ptr<Module>> submodules_;
+    std::unordered_map<std::string, Parameter> parameters_;
+};
+} // namespace infinicore::nn
\ No newline at end of file

include/infinicore/nn/parameter.hpp

+#pragma once
+
+#include "../tensor.hpp"
+
+namespace infinicore::nn {
+class Parameter : public Tensor {
+public:
+    Parameter(const Shape &shape,
+              const DataType &dtype,
+              const Device &device);
+
+    void load_blob(const void *data);
+};
+} // namespace infinicore::nn

include/infinicore/tensor.hpp

@@ -110,6 +110,10 @@ public:
 
     Size size(size_t dim) const;
 
+    size_t element_size() const;
+
+    size_t nbytes() const;
+
     Stride stride(size_t dim) const;
 
     DataType dtype() const;
@@ -142,7 +146,6 @@ public:
 
     /**
      * Copy Data from another tensor to this tensor.
-     * Currently, only contigous tensors of the same dtype and shape are supported.
      *
      * @param src The source tensor to copy from
      *

src/infinicore/nn/module.cc

+#include "infinicore/nn/module.hpp"
+
+namespace infinicore::nn {
+const std::unordered_map<std::string, Parameter> &Module::state_dict() const {
+    return parameters_;
+}
+
+void Module::load_state_dict(const std::unordered_map<std::string, Tensor> &_state_dict) {
+    for (auto &p : parameters_) {
+        load_parameter(p.first, p.second);
+    }
+}
+
+void Module::load_parameter(const std::string &name, const Tensor &param) {
+    parameters_[name]->copy_from(param);
+}
+
+void Module::load_parameter_from_blob(const std::string &name, const void *data) {
+    auto param = parameters_[name];
+    param.load_blob(data);
+}
+
+Tensor Module::register_parameter(const std::string &name, Parameter param) {
+    parameters_[name] = param;
+    return param;
+}
+
+} // namespace infinicore::nn

src/infinicore/nn/parameter.cc

+#include "infinicore/nn/parameter.hpp"
+
+#include "infinicore/context/context.hpp"
+
+#include <cstring>
+
+namespace infinicore::nn {
+Parameter::Parameter(
+    const Shape &shape,
+    const DataType &dtype,
+    const Device &device)
+    : Tensor(Tensor::empty(shape, dtype, device, false)) {
+}
+
+void Parameter::load_blob(const void *data) {
+    auto buffer = Tensor::empty(impl_->shape(), impl_->dtype(), Device(Device::Type::CPU, 0), true);
+    std::memcpy(buffer->data(), data, buffer->nbytes());
+    infinicore::context::memcpyH2D(impl_->data(), buffer->data(), buffer->nbytes());
+    infinicore::context::syncStream();
+}
+} // namespace infinicore::nn

src/infinicore/tensor/tensor.cc

@@ -117,6 +117,14 @@ Size TensorImpl::numel() const {
     return total;
 }
 
+size_t TensorImpl::element_size() const {
+    return dsize(dtype());
+}
+
+size_t TensorImpl::nbytes() const {
+    return numel() * element_size();
+}
+
 Size TensorImpl::size(size_t dim) const {
     return meta_.shape[dim];
 }