layer.h

Go to the documentation of this file.

#pragma once
#include <vector>
#include <string>
#include "compute/operation.h"

namespace magmadnn {
namespace layer {

template <typename T>
class Layer {
public:
    
    virtual std::vector<op::Operation<T> *> get_weights() = 0;

    virtual op::Operation<T>* out() {
        return output;
    }

    op::Operation<T> *get_input() { return input; }
    op::Operation<T> *get_output() { return output; }

    std::vector<unsigned int> get_input_shape() const { return input_shape; }
    std::vector<unsigned int> get_output_shape() const { return this->output->get_output_shape(); }

    unsigned int get_input_shape(unsigned int i) const {
        assert( i < input_shape.size() );
        return input_shape[i];
    }

    unsigned int get_output_shape(unsigned int i) const {
        assert( i < output_shape.size() );
        return output_shape[i];
    }

    void set_name(std::string name) { this->name = name; }

    std::string get_name() const { return this->name; }

protected:
    Layer(std::vector<unsigned int> input_shape, op::Operation<T> *input) : 
        input_shape(input_shape), input(input) {}

    std::vector<unsigned int> input_shape;
    std::vector<unsigned int> output_shape;

    op::Operation<T> *input;
    op::Operation<T> *output;

    std::string name;

};

}   // namespace layer
}   // namespace magmadnn