Migrate batch norm layer

This commit is contained in:
2025-11-21 23:24:14 +01:00
parent 5679dc0a50
commit fd4775faa4
11 changed files with 181 additions and 364 deletions

View File

@@ -2,7 +2,7 @@
#include "kernels/activation_functions.cuh"
#include "kernels/convolution.cuh"
#include "kernels/matmul.cuh"
#include "kernels/pooling.cuh"
#include "kernels/pool.cuh"
#include "utils/cuda_helper.cuh"
using namespace CUDANet::Backend;
@@ -112,7 +112,7 @@ CUDANet::Tensor& CUDA::conv2d(
return output;
}
CUDANet::Tensor& CUDA::maxPool2d(
CUDANet::Tensor& CUDA::max_pool2d(
const CUDANet::Tensor& input,
CUDANet::Tensor& output,
CUDANet::Shape input_shape,
@@ -138,7 +138,7 @@ CUDANet::Tensor& CUDA::maxPool2d(
return output;
}
CUDANet::Tensor& CUDA::avgPool2d(
CUDANet::Tensor& CUDA::avg_pool2d(
const CUDANet::Tensor& input,
CUDANet::Tensor& output,
CUDANet::Shape input_shape,
@@ -162,4 +162,53 @@ CUDANet::Tensor& CUDA::avgPool2d(
CUDA_CHECK(cudaDeviceSynchronize());
return output;
}
CUDANet::Tensor& CUDA::batch_norm(
const CUDANet::Tensor& input,
CUDANet::Tensor& output,
CUDANet::Shape input_shape,
CUDANet::Tensor& weights,
CUDANet::Tensor& biases,
CUDANet::Tensor& running_mean,
CUDANet::Tensor& running_var,
CUDANet::Tensor& epsilon
) {
auto gridSize =
(input_shape[0] * input_shape[1] + BLOCK_SIZE - 1) / BLOCK_SIZE;
for (int i = 0; i < input_shape[2]; i++) {
// Subtract mean from input
Kernels::vec_scalar_sub<<<gridSize, BLOCK_SIZE>>>(
input.data<float>() + i * input_shape[0] * input_shape[1],
output.data<float>() + i * input_shape[0] * input_shape[1],
&running_mean.data<float>()[i], input_shape[0] * input_shape[1]
);
CUDA_CHECK(cudaGetLastError());
// Divide by sqrt(running_var + epsilon)
Kernels::vec_scale<<<gridSize, BLOCK_SIZE>>>(
output.data<float>() + i * input_shape[0] * input_shape[1],
output.data<float>() + i * input_shape[0] * input_shape[1],
&running_var.data<float>()[i], epsilon.data<float>(), input_shape[0] * input_shape[1]
);
CUDA_CHECK(cudaGetLastError());
// Multiply by weights
Kernels::vec_scalar_mul<<<gridSize, BLOCK_SIZE>>>(
output.data<float>() + i * input_shape[0] * input_shape[1],
output.data<float>() + i * input_shape[0] * input_shape[1], &weights.data<float>()[i],
input_shape[0] * input_shape[1]
);
CUDA_CHECK(cudaGetLastError());
// Add biases
Kernels::vec_scalar_add<<<gridSize, BLOCK_SIZE>>>(
output.data<float>() + i * input_shape[0] * input_shape[1],
output.data<float>() + i * input_shape[0] * input_shape[1], &biases.data<float>()[i],
input_shape[0] * input_shape[1]
);
CUDA_CHECK(cudaGetLastError());
}
}

View File

@@ -1,120 +0,0 @@
#include <vector>
#include "activation.hpp"
#include "batch_norm.hpp"
#include "cuda_helper.cuh"
#include "layer.hpp"
#include "matmul.cuh"
#include "vector.cuh"
using namespace CUDANet::Layers;
void BatchNorm2d::initCUDA() {
d_output = nullptr;
CUDA_CHECK(cudaMalloc(
(void **)&d_output,
sizeof(float) * inputSize.first * inputSize.second * inputChannels
));
d_running_mean = nullptr;
CUDA_CHECK(
cudaMalloc((void **)&d_running_mean, sizeof(float) * inputChannels)
);
d_running_var = nullptr;
CUDA_CHECK(
cudaMalloc((void **)&d_running_var, sizeof(float) * inputChannels)
);
d_weights = nullptr;
CUDA_CHECK(cudaMalloc((void **)&d_weights, sizeof(float) * inputChannels));
d_biases = nullptr;
CUDA_CHECK(cudaMalloc((void **)&d_biases, sizeof(float) * inputChannels));
d_length = nullptr;
float length = (float)inputSize.first * inputSize.second;
CUDA_CHECK(cudaMalloc((void **)&d_length, sizeof(float)));
CUDA_CHECK(
cudaMemcpy(d_length, &length, sizeof(float), cudaMemcpyHostToDevice)
);
d_epsilon = nullptr;
CUDA_CHECK(cudaMalloc((void **)&d_epsilon, sizeof(float)));
CUDA_CHECK(
cudaMemcpy(d_epsilon, &epsilon, sizeof(float), cudaMemcpyHostToDevice)
);
gridSize =
(inputSize.first * inputSize.second + BLOCK_SIZE - 1) / BLOCK_SIZE;
}
void BatchNorm2d::delCUDA() {
cudaFree(d_output);
cudaFree(d_running_mean);
cudaFree(d_running_var);
cudaFree(d_weights);
cudaFree(d_biases);
cudaFree(d_length);
cudaFree(d_epsilon);
}
void BatchNorm2d::toCuda() {
CUDA_CHECK(cudaMemcpy(
d_weights, weights.data(), sizeof(float) * inputChannels,
cudaMemcpyHostToDevice
));
CUDA_CHECK(cudaMemcpy(
d_biases, biases.data(), sizeof(float) * inputChannels,
cudaMemcpyHostToDevice
));
CUDA_CHECK(cudaMemcpy(
d_running_mean, running_mean.data(), sizeof(float) * inputChannels,
cudaMemcpyHostToDevice
));
CUDA_CHECK(cudaMemcpy(
d_running_var, running_var.data(), sizeof(float) * inputChannels,
cudaMemcpyHostToDevice
));
}
float *BatchNorm2d::forwardCUDA(const float *d_input) {
// Compute per-channel batch normalization
for (int i = 0; i < inputChannels; i++) {
// Subtract mean from input
Kernels::vec_scalar_sub<<<gridSize, BLOCK_SIZE>>>(
d_input + i * inputSize.first * inputSize.second,
d_output + i * inputSize.first * inputSize.second,
&d_running_mean[i], inputSize.first * inputSize.second
);
CUDA_CHECK(cudaGetLastError());
// Divide by sqrt(running_var + epsilon)
Kernels::vec_scale<<<gridSize, BLOCK_SIZE>>>(
d_output + i * inputSize.first * inputSize.second,
d_output + i * inputSize.first * inputSize.second,
&d_running_var[i], d_epsilon, inputSize.first * inputSize.second
);
CUDA_CHECK(cudaGetLastError());
// Multiply by weights
Kernels::vec_scalar_mul<<<gridSize, BLOCK_SIZE>>>(
d_output + i * inputSize.first * inputSize.second,
d_output + i * inputSize.first * inputSize.second, &d_weights[i],
inputSize.first * inputSize.second
);
CUDA_CHECK(cudaGetLastError());
// Add biases
Kernels::vec_scalar_add<<<gridSize, BLOCK_SIZE>>>(
d_output + i * inputSize.first * inputSize.second,
d_output + i * inputSize.first * inputSize.second, &d_biases[i],
inputSize.first * inputSize.second
);
CUDA_CHECK(cudaGetLastError());
}
activation->activate(d_output);
return d_output;
}

View File

@@ -23,7 +23,12 @@ void CUDA::print(const CUDANet::Tensor &input) {
}
void CUDA::zero(CUDANet::Tensor &input) {
CUDA_CHECK(cudaMemset(input.data<float>(), 0, sizeof(float) * input.numel()));
fill(input, 0);
}
void CUDA::fill(CUDANet::Tensor &input, int value) {
CUDA_CHECK(cudaMemset(input.data<float>(), value, sizeof(float) * input.numel()));
}
void CUDA::copy_to_device(CUDANet::Tensor &tensor, void *data, size_t size) {

View File

@@ -51,7 +51,7 @@ AvgPool2d::~AvgPool2d() {}
CUDANet::Tensor& AvgPool2d::forward(CUDANet::Tensor& input) {
output.zero();
backend->avgPool2d(
backend->avg_pool2d(
input,
output,
in_shape,

View File

@@ -9,125 +9,95 @@
using namespace CUDANet::Layers;
BatchNorm2d::BatchNorm2d(
shape2d inputSize,
int inputChannels,
float epsilon,
ActivationType activationType
CUDANet::Shape input_shape,
float eps,
CUDANet::Backend *backend
)
: inputSize(inputSize), inputChannels(inputChannels), epsilon(epsilon) {
activation = new Activation(
activationType, inputSize.first * inputSize.second * inputChannels
: in_shape(input_shape), backend(backend) {
if (in_shape.size() != 3) {
throw InvalidShapeException("input", 3, in_shape.size());
}
epsilon = CUDANet::Tensor({1}, CUDANet::DType::FLOAT32, backend);
epsilon.set_data<float>(&eps);
running_mean = CUDANet::Tensor({in_shape[2]}, CUDANet::DType::FLOAT32, backend);
running_mean.zero();
running_var = CUDANet::Tensor({in_shape[2]}, CUDANet::DType::FLOAT32, backend);
running_var.fill(1);
weights = CUDANet::Tensor({in_shape[2]}, CUDANet::DType::FLOAT32, backend);
weights.fill(1);
biases = CUDANet::Tensor({in_shape[2]}, CUDANet::DType::FLOAT32, backend);
biases.zero();
output = CUDANet::Tensor(in_shape, CUDANet::DType::FLOAT32, backend);
}
BatchNorm2d::~BatchNorm2d() {}
CUDANet::Tensor& BatchNorm2d::forward(CUDANet::Tensor& input) {
output.zero();
backend->batch_norm(
input,
output,
in_shape,
weights,
biases,
running_mean,
running_var,
epsilon
);
weights.resize(inputChannels);
biases.resize(inputChannels);
running_mean.resize(inputChannels);
running_var.resize(inputChannels);
initializeWeights();
initializeBiases();
initializeRunningMean();
initializeRunningVar();
#ifdef USE_CUDA
initCUDA();
toCuda();
#endif
return output;
}
BatchNorm2d::~BatchNorm2d() {
#ifdef USE_CUDA
delCUDA();
#endif
CUDANet::Shape BatchNorm2d::input_shape() {
return in_shape;
}
void BatchNorm2d::initializeWeights() {
std::fill(weights.begin(), weights.end(), 1.0f);
CUDANet::Shape BatchNorm2d::output_shape() {
return in_shape;
}
void BatchNorm2d::initializeBiases() {
std::fill(biases.begin(), biases.end(), 0.0f);
size_t BatchNorm2d::input_size() {
return sizeof(float) * in_shape[0] * in_shape[1] * in_shape[2];
}
void BatchNorm2d::initializeRunningMean() {
std::fill(running_mean.begin(), running_mean.end(), 0.0f);
size_t BatchNorm2d::output_size() {
return sizeof(float) * in_shape[0] * in_shape[1] * in_shape[2];
}
void BatchNorm2d::initializeRunningVar() {
std::fill(running_var.begin(), running_var.end(), 1.0f);
void BatchNorm2d::set_weights(void* input) {
weights.set_data<float>(static_cast<float*>(input));
}
void BatchNorm2d::setWeights(const float* weights_input) {
std::copy(weights_input, weights_input + weights.size(), weights.begin());
#ifdef USE_CUDA
toCuda();
#endif
}
std::vector<float> BatchNorm2d::getWeights() {
CUDANet::Tensor& BatchNorm2d::get_weights() {
return weights;
}
void BatchNorm2d::setBiases(const float* biases_input) {
std::copy(biases_input, biases_input + biases.size(), biases.begin());
#ifdef USE_CUDA
toCuda();
#endif
void BatchNorm2d::set_biases(void* input) {
biases.set_data<float>(static_cast<float*>(input));
}
std::vector<float> BatchNorm2d::getBiases() {
CUDANet::Tensor& BatchNorm2d::get_biases() {
return biases;
}
void BatchNorm2d::setRunningMean(const float* running_mean_input) {
std::copy(
running_mean_input, running_mean_input + inputChannels,
running_mean.begin()
);
#ifdef USE_CUDA
toCuda();
#endif
void BatchNorm2d::set_running_mean(void* input) {
running_mean.set_data<float>(static_cast<float*>(input));
}
std::vector<float> BatchNorm2d::getRunningMean() {
CUDANet::Tensor& BatchNorm2d::get_running_mean() {
return running_mean;
}
void BatchNorm2d::setRunningVar(const float* running_var_input) {
std::copy(
running_var_input, running_var_input + inputChannels,
running_var.begin()
);
#ifdef USE_CUDA
toCuda();
#endif
void BatchNorm2d::set_running_var(void* input) {
running_var.set_data<float>(static_cast<float*>(input));
}
std::vector<float> BatchNorm2d::getRunningVar() {
CUDANet::Tensor& BatchNorm2d::get_running_var() {
return running_var;
}
int BatchNorm2d::getInputSize() {
return inputSize.first * inputSize.second * inputChannels;
}
int BatchNorm2d::getOutputSize() {
return inputSize.first * inputSize.second * inputChannels;
}
shape2d BatchNorm2d::getOutputDims() {
return inputSize;
}
float* BatchNorm2d::forwardCPU(const float* input) {
throw std::logic_error("Not implemented");
}
float* BatchNorm2d::forward(const float* input) {
#ifdef USE_CUDA
return forwardCUDA(input);
#else
return forwardCPU(input);
#endif
}

View File

@@ -47,7 +47,7 @@ Conv2d::Conv2d(
};
output = CUDANet::Tensor(
Shape{out_shape[0] * out_shape[1] * out_shape[3]},
Shape{out_shape[0], out_shape[1], out_shape[3]},
CUDANet::DType::FLOAT32, backend
);

View File

@@ -50,7 +50,7 @@ MaxPool2d::~MaxPool2d() {}
CUDANet::Tensor& MaxPool2d::forward(CUDANet::Tensor& input) {
output.zero();
backend->maxPool2d(
backend->max_pool2d(
input, output, in_shape, pool_shape, stride_shape, padding_shape,
out_shape
);