mirror of
https://github.com/lordmathis/CUDANet.git
synced 2025-12-22 22:34:22 +00:00
55 lines
1.1 KiB
C++
55 lines
1.1 KiB
C++
#pragma once
|
|
|
|
#include "layer.hpp"
|
|
|
|
namespace CUDANet::Layers {
|
|
|
|
class BatchNorm2d : public Layer {
|
|
public:
|
|
BatchNorm2d(CUDANet::Shape input_shape, float epsilon, CUDANet::Backend *backend);
|
|
|
|
~BatchNorm2d();
|
|
|
|
CUDANet::Tensor& forward(CUDANet::Tensor& input) override;
|
|
|
|
CUDANet::Shape input_shape() override;
|
|
|
|
CUDANet::Shape output_shape() override;
|
|
|
|
size_t input_size() override;
|
|
|
|
size_t output_size() override;
|
|
|
|
void set_weights(void* input) override;
|
|
|
|
CUDANet::Tensor& get_weights() override;
|
|
|
|
void set_biases(void* input) override;
|
|
|
|
CUDANet::Tensor& get_biases() override;
|
|
|
|
void set_running_mean(void* input);
|
|
|
|
CUDANet::Tensor& get_running_mean();
|
|
|
|
void set_running_var(void* input);
|
|
|
|
CUDANet::Tensor& get_running_var();
|
|
|
|
private:
|
|
CUDANet::Shape in_shape;
|
|
CUDANet::Tensor epsilon;
|
|
|
|
CUDANet::Tensor running_mean;
|
|
CUDANet::Tensor running_var;
|
|
|
|
CUDANet::Tensor weights;
|
|
CUDANet::Tensor biases;
|
|
|
|
CUDANet::Tensor output;
|
|
|
|
CUDANet::Backend *backend;
|
|
};
|
|
|
|
} // namespace CUDANet::Layers
|