mnist_classifier_example.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <vector>
 
#include "pulsatrix/adam_optimizer.hpp"
#include "pulsatrix/conv2d_module.hpp"
#include "pulsatrix/cross_entropy_loss.hpp"
#include "pulsatrix/flatten_module.hpp"
#include "pulsatrix/linear_module.hpp"
#include "pulsatrix/metrics_sink.hpp"
#include "pulsatrix/relu_module.hpp"
 
namespace pulsatrix {
 
class MnistConvNet {
public:
    explicit MnistConvNet(DeviceBackend* backend, unsigned seed = 42);
 
    [[nodiscard]] Tensor forward(const Tensor& image);
 
    [[nodiscard]] int64_t predict(const Tensor& image);
 
    float train_step(const Tensor& image, int64_t target_class, AdamOptimizer& optimizer, MetricsSink& sink,
                      int step);
 
    [[nodiscard]] const Tensor& classifier_weight() const { return classifier_.weight(); }
 
    [[nodiscard]] std::vector<Module*> modules() { return {&conv_, &relu_, &flatten_, &classifier_}; }
 
private:
    DeviceBackend* backend_;
    Conv2DModule conv_;
    ReluModule relu_;
    FlattenModule flatten_;
    LinearModule classifier_;
    CrossEntropyLoss loss_;
};
 
}  // namespace pulsatrix