model-io.h

Go to the documentation of this file.

// Copyright (c) 2021, Aalto University, developed by Erik Schultheis
// All rights reserved.
//
// SPDX-License-Identifier: MIT
 
#ifndef DISMEC_MODEL_IO_H
#define DISMEC_MODEL_IO_H
 
#include <Eigen/Dense>
#include <vector>
#include <memory>
#include <filesystem>
#include <optional>
#include <future>
#include "fwd.h"
#include "data/types.h"
#include <boost/iterator/iterator_adaptor.hpp>
 
namespace dismec::io
{
 
    namespace model
    {
        using dismec::model::Model;
        using std::filesystem::path;
 
        enum class WeightFormat {
            DENSE_TXT  = 0,      
            SPARSE_TXT = 1,      
            DENSE_NPY  = 2,      
            NULL_FORMAT = 3      
        };
 
        WeightFormat parse_weights_format(std::string_view name);
        const char* to_string(WeightFormat format);
 
        struct SaveOption {
            int Precision  = 6;     
            double Culling = 0;     
            int SplitFiles = 4096;  
            WeightFormat Format = WeightFormat::DENSE_TXT;      
        };
 
        // \todo should we by default overwrite files, or refuse?
        void save_model(const path& target_file, const std::shared_ptr<const Model>& model, SaveOption options);
 
 
        std::shared_ptr<Model> load_model(path source);
 
        struct WeightFileEntry {
            label_id_t First;
            long Count;
            std::string FileName;
            WeightFormat Format;
        };
 
 
        class PartialModelIO {
        public:
            [[nodiscard]] long num_labels() const noexcept { return m_TotalLabels; }
 
            [[nodiscard]] long num_features() const noexcept { return m_NumFeatures; }
 
        protected:
            PartialModelIO() = default;
            ~PartialModelIO() = default;
 
            void read_metadata_file(const path& meta_file);
 
            long m_TotalLabels = -1;
            long m_NumFeatures = -1;
 
            std::vector<WeightFileEntry> m_SubFiles;
            using weight_file_iter_t = std::vector<WeightFileEntry>::const_iterator;
 
            void insert_sub_file(const WeightFileEntry& data);
 
            [[nodiscard]] weight_file_iter_t label_lower_bound(label_id_t pos) const;
        };
 
        class PartialModelSaver : public PartialModelIO {
        public:
            PartialModelSaver(path target_file, SaveOption options, bool load_partial=false);
 
            std::future<WeightFileEntry> add_model(const std::shared_ptr<const Model>& model,
                                                   const std::optional<std::string>& file_path={});
 
            using PartialModelIO::insert_sub_file;
 
            void update_meta_file();
 
            void finalize();
 
             [[nodiscard]] std::pair<label_id_t, label_id_t> get_missing_weights() const;
 
              [[nodiscard]] bool any_weight_vector_for_interval(label_id_t begin, label_id_t end) const;
 
        private:
            SaveOption m_Options;
            path m_MetaFileName;
        };
 
        class PartialModelLoader : public PartialModelIO {
        public:
            enum ESparseMode {
                DEFAULT,
                FORCE_SPARSE,
                FORCE_DENSE
            };
 
            explicit PartialModelLoader(path meta_file, ESparseMode mode=DEFAULT);
 
            const path& meta_file_path() const { return m_MetaFileName; }
 
            [[nodiscard]] std::shared_ptr<Model> load_model(label_id_t label_begin, label_id_t label_end) const;
 
            [[nodiscard]] std::shared_ptr<Model> load_model(int index) const;
 
            bool validate() const;
 
            [[nodiscard]] long num_weight_files() const;
 
            struct SubModelRangeSpec {
                weight_file_iter_t FilesBegin;
                weight_file_iter_t FilesEnd;
                label_id_t LabelsBegin{0};
                label_id_t LabelsEnd{0};
            };
 
            [[nodiscard]] SubModelRangeSpec get_loading_range(label_id_t label_begin, label_id_t label_end) const;
        private:
            path m_MetaFileName;
            ESparseMode m_SparseMode;
        };
    }
 
    using model::WeightFormat;
    using model::SaveOption;
    using model::save_model;
    using model::load_model;
    using model::PartialModelSaver;
    using model::PartialModelLoader;
}
 
#endif //DISMEC_MODEL_IO_H