dnf_torch.h

 
#ifndef _DNF_H
#define _DNF_H
 
#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <assert.h>
#include <torch/torch.h>
#include <thread>
#include <iostream>
#include <deque>
 
#ifdef NDEBUG
constexpr bool debugOutput = false;
#else
constexpr bool debugOutput = true;
#endif
 
class DNF
{
public:
    enum ActMethod
    {
        Act_Sigmoid = 1,
        Act_Tanh = 2,
        Act_ReLU = 3,
        Act_NONE = 0
    };
 
private:
    struct Net : public torch::nn::Module
    {
        std::vector<torch::nn::Linear> fc;
        Net(int nLayers, int nInput, bool withBias = false);
        torch::Tensor forward(torch::Tensor x, ActMethod am);
    };
 
public:
    DNF(const int nLayers,
        const int nTaps,
        const ActMethod am = Act_Tanh,
        const bool tryGPU = false);
 
    void setLearningRate(float mu);
 
    float filter(const float signal, const float noise);
 
    inline int getSignalDelaySteps() const
    {
        return signalDelayLineLength;
    }
 
    inline float getDelayedSignal() const
    {
        return signal_delayLine.get(0);
    }
 
    inline float getRemover() const
    {
        return remover;
    }
 
    inline float getOutput() const
    {
        return f_nn;
    }
 
    const std::vector<float> getLayerWeightDistances() const;
 
    float getWeightDistance() const;
 
    const torch::Device getTorchDevice() const
    {
        return device;
    }
 
    const Net getModel() const
    {
        return model;
    }
 
    static constexpr double xavierGain = 0.01;
 
private:
    class DelayLine
    {
    public:
        void init(int delay)
        {
            delaySamples = delay;
            buffer = std::deque<float>(delaySamples, 0.0f);
        }
 
        inline float process(float input)
        {
            float output = buffer.front();
            buffer.pop_front();
            buffer.push_back(input);
            return output;
        }
 
        float get(int i) const
        {
            return buffer[i];
        }
 
        float getNewest() const
        {
            return buffer.back();
        }
 
    private:
        int delaySamples = 0;
        std::deque<float> buffer;
    };
 
    void saveInitialParameters()
    {
        for (const auto &p : model.parameters())
        {
            initialParameters.push_back(p.detach().clone());
        }
    }
 
    const int noiseDelayLineLength;
    const int signalDelayLineLength;
    const ActMethod actMethod;
    Net model;
    torch::optim::SGD optimizer;
    std::vector<torch::Tensor> initialParameters;
    DelayLine signal_delayLine;
    DelayLine noise_delayLine;
    float remover = 0;
    float f_nn = 0;
    torch::Device device = torch::kCPU;
};
 
#endif