src/nn/module.cpp - master - vuron/adept

src/nn/module.cpp
94 строки2 KB
kolkir
Fix training issues
06 мар 2025, 23:54
06 мар 2025, 23:54efe26e8
#include <adept/nn/module.hpp>

namespace adept {

Module::Module(const std::string& name) : name_(name) {}

const std::string& Module::name() const noexcept {
  return name_;
}

void Module::register_module(std::string name, std::shared_ptr<Module> module) {
  if (module.get() != nullptr)
    children_.emplace(std::move(name), std::move(module));
}

void Module::register_parameter(std::string name, Variable parameter) {
  // some parameters like biases can be undefined
  if (parameter.defined()) {
    parameters_.emplace(std::move(name), std::move(parameter));
  }
}

std::vector<Variable> Module::parameters() const {
  std::vector<Variable> parameters;
  for (auto& param_item : parameters_) {
    parameters.push_back(param_item.second);
  }

  for (auto& child_item : children_) {
    auto params = child_item.second->parameters();
    parameters.insert(parameters.end(), std::make_move_iterator(std::begin(params)),
                      std::make_move_iterator(std::end(params)));
  }
  return parameters;
}

namespace {
template <typename... Args>
std::string join_name(const std::string& arg, const Args&... args) {
  std::ostringstream oss;
  oss << arg;
  if constexpr (sizeof...(args) > 0) {
    if (!arg.empty())
      oss << "_";
    oss << join_name(args...);
  }
  return oss.str();
}
}  // namespace

void Module::save(OutputSerializer& output, const std::string& prefix) const {
  for (auto& param : parameters_) {
    output.write(join_name(prefix, name_, param.first), param.second);
  }
  for (auto& child_item : children_) {
    child_item.second->save(output, join_name(name_, child_item.first));
  }
}

void Module::load(InputSerializer& input, const std::string& prefix) {
  for (auto& param : parameters_) {
    input.read(join_name(prefix, name_, param.first), param.second);
  }
  for (auto& child_item : children_) {
    child_item.second->load(input, join_name(name_, child_item.first));
  }
}

void Module::train() {
  is_train_mode_ = true;
  for (auto& child_item : children_) {
    child_item.second->train();
  }
}
void Module::eval() {
  is_train_mode_ = false;
  for (auto& child_item : children_) {
    child_item.second->eval();
  }
}

void Module::to_string(std::ostream& stream) const {
  stream << name_ << ": [\n";
  for (auto& param : parameters_) {
    stream << param.first << ":\n";
    stream << param.second << "\n";
  }
  for (auto& child_item : children_) {
    child_item.second->to_string(stream);
  }
  stream << "]\n";
}

}  // namespace adept