network_wrapper.hpp

Go to the documentation of this file.

#pragma once
 
#include "layer.hpp"   // CPU layer tags
#include "network.hpp" // CPU implementation
 
#ifdef __CUDACC__
  #include "cuda/kernels.cuh" // CUDA Enums
  #include "cuda/network.cuh" // GPU Implementation
#endif
 
#include <iostream>
#include <memory>
 
struct CpuBackend {};
struct CudaBackend {};
 
template <typename T> struct ActivationToEnum;
 
template <> struct ActivationToEnum<cpu_mlp::Linear> {
#ifdef __CUDACC__
  static constexpr cuda_mlp::ActivationType value = cuda_mlp::ActivationType::Linear;
#endif
};
template <> struct ActivationToEnum<cpu_mlp::Sigmoid> {
#ifdef __CUDACC__
  static constexpr cuda_mlp::ActivationType value = cuda_mlp::ActivationType::Sigmoid;
#endif
};
template <> struct ActivationToEnum<cpu_mlp::Tanh> {
#ifdef __CUDACC__
  static constexpr cuda_mlp::ActivationType value = cuda_mlp::ActivationType::Tanh;
#endif
};
template <> struct ActivationToEnum<cpu_mlp::ReLU> {
#ifdef __CUDACC__
  static constexpr cuda_mlp::ActivationType value = cuda_mlp::ActivationType::ReLU;
#endif
};
 
template <typename Backend> class NetworkWrapper;
 
template <> class NetworkWrapper<CpuBackend> {
public:
  using InternalNetwork = cpu_mlp::Network;
 
  NetworkWrapper() = default;
 
  template <int In, int Out, typename Activation> void addLayer() { network_.addLayer<In, Out, Activation>(); }
 
  void bindParams() { network_.bindParams(); }
  void bindParams(unsigned int seed) { network_.bindParams(seed); }
 
  InternalNetwork &getInternal() { return network_; }
  const InternalNetwork &getInternal() const { return network_; }
 
  size_t getParamsSize() const { return network_.getSize(); }
 
private:
  InternalNetwork network_;
};
 
#ifdef __CUDACC__
template <> class NetworkWrapper<CudaBackend> {
public:
  using InternalNetwork = cuda_mlp::CudaNetwork;
 
  explicit NetworkWrapper(cuda_mlp::CublasHandle &handle) : network_(handle) {}
 
  template <int In, int Out, typename Activation> void addLayer() {
    network_.addLayer(In, Out, ActivationToEnum<Activation>::value);
  }
 
  void bindParams() { network_.bindParams(); }
  void bindParams(unsigned int seed) { network_.bindParams(seed); }
 
  InternalNetwork &getInternal() { return network_; }
  const InternalNetwork &getInternal() const { return network_; }
 
  size_t getParamsSize() const { return network_.params_size(); }
 
private:
  InternalNetwork network_;
};
#endif