net.hpp

#pragma once
#include <dlprim/operator.hpp>
#include <map>
#include <vector>
#include <list>

namespace dlprim {
    class SharedResource;
    class ModelBase;

    class Net {
    public:

        Net(Context &ctx);

        ~Net() {}
        Net const &operator=(Net const &) = delete;
        Net(Net const &) = delete;
        
        Net &operator=(Net &&) = default;
        Net(Net &&) = default;

        void load_model(ModelBase &model);
        void load_from_json(json::value const &v);
        void load_from_json_file(std::string const &name);

        void add_input_tensor(std::string const &name,TensorSpecs const &ts,bool requires_gradient=false);
        void mark_output_tensor(std::string const &name);
        void set_loss_weight(std::string const &name,float loss_weight=1.0);

        void load_parameters(std::istream &s,bool allow_missing=false);
        void load_parameters(std::string const &name,bool allow_missing=false);
        void load_parameters(ModelBase &model,bool allow_missing=false);

        void save_parameters(std::string const &fname);

        void load_parameters_from_hdf5(std::string const &fname,bool allow_missing=false);
        void save_parameters_to_hdf5(std::string const &fname);

        bool keep_intermediate_tensors() const
        {
            return keep_intermediate_tensors_;
        }
        void keep_intermediate_tensors(bool keep) 
        {
            keep_intermediate_tensors_ = keep;
        }

        void add_operator(  std::unique_ptr<Operator> op,
                            std::string const &name,
                            std::vector<std::string> const &inputs,
                            std::vector<std::string> const &outputs,
                            std::vector<std::string> const &parameters = std::vector<std::string>(),
                            bool frozen=false);

        void mode(CalculationsMode mode); 

        CalculationsMode mode() const
        {
            return mode_;
        }

        std::shared_ptr<SharedResource> shared_resource()
        {
            return shared_resource_;
        }

        void setup();
        void reshape();

        void initialize_parameters(ExecutionContext const &e);

        void copy_parameters_to_device(); 
        void copy_parameters_to_host(); 
        void clear_memory();

        std::map<std::string,Tensor> &tensors()
        {
            return tensors_;
        }
        std::map<std::string,Tensor> &tensor_diffs()
        {
            return tensors_diff_;
        }

        std::map<std::string,Tensor> &params()
        {
            return parameters_;
        }
        std::map<std::string,Tensor> &param_diffs()
        {
            return parameters_diff_;
        }

        Tensor &tensor(std::string const &name)
        {
            auto p=tensors_.find(name);
            if(p == tensors_.end())
                throw ValidationError("Unknown tensor name:" + name);
            return p->second;
        }
        Tensor &tensor_diff(std::string const &name)
        {
            auto p=tensors_diff_.find(name);
            if(p == tensors_diff_.end())
                throw ValidationError("Unknown tensor name:" + name);
            return p->second;
        }
        
        Tensor &param(std::string const &name)
        {
            auto p=parameters_.find(name);
            if(p == parameters_.end())
                throw ValidationError("Unknown parameter name:" + name);
            return p->second;
        }
        Tensor &param_diff(std::string const &name)
        {
            auto p=parameters_diff_.find(name);
            if(p == parameters_diff_.end())
                throw ValidationError("Unknown parameter name:" + name);
            return p->second;
        }

        void forward(ExecutionContext const &ectx,bool sync=false);
        void backward(ExecutionContext const &ectx,bool sync=false);

        std::vector<std::string> const &input_names()
        {
            return inputs_;
        }

        std::vector<std::string> const &output_names()
        {
            return outputs_;
        }

        Tensor &input(unsigned id)
        {
            if(id >= inputs_.size())
                throw ValidationError("Invalid input id");
            return tensor(inputs_[id]);
        }

        Tensor &output(unsigned id)
        {
            if(id >= outputs_.size())
                throw ValidationError("Invalid output id");
            return tensor(outputs_[id]);
        }

    private:
        struct Connection {
            std::string name;
            std::unique_ptr<Operator> op;
            
            std::vector<std::string> parameter_names;
            std::vector<std::string> input_names;
            std::vector<Tensor>      input_tensors;
            std::vector<TensorSpecs> input_specs;

            std::vector<std::string> output_names;
            std::vector<Tensor>      output_tensors;
            std::vector<TensorSpecs> output_specs;
            
            std::vector<TensorSpecs> parameter_specs;
            std::vector<Tensor>      parameters;

            std::vector<TensorAndGradient> in_grad;
            std::vector<TensorAndGradient> out_grad;
            std::vector<TensorAndGradient> param_grad;
            size_t ws_size;
            int gradient_flags;
            bool frozen;
        };

        void setup_ws();
        void mark_backpropagating_edges();
        void allocate_tensors();
        void allocate_aliases();
        bool is_loss(std::string const &name);
        void allocate_optimized_chunks(bool forward_only);
        void tensor_use_list(std::vector<std::list<std::string> > &start,
                              std::vector<std::list<std::string> > &stop);
        void allocate_chunks();
        void load_header(std::istream &f,json::value &v);


        Context ctx_;
        std::map<std::string,TensorSpecs> tensor_specs_;
        std::map<std::string,TensorSpecs> parameter_specs_;

        std::vector<Connection> connections_;
        std::map<std::string,unsigned> connections_index_;

        Tensor workspace_;
        std::map<std::string,std::string> alias_sources_;
        std::map<std::string,Tensor> tensors_;
        std::map<std::string,Tensor> tensors_diff_;

        std::map<std::string,Tensor> parameters_;
        std::map<std::string,Tensor> parameters_diff_;
        std::map<std::string,float> loss_weights_;
        std::vector<std::string> inputs_;
        std::vector<std::string> outputs_;
        std::vector<Tensor> memory_;

        std::shared_ptr<SharedResource> shared_resource_;

        CalculationsMode mode_;
        bool keep_intermediate_tensors_;
    };
};