diff --git a/src/batchnorm_layer.c b/src/batchnorm_layer.c index a025fc4a..a2a850ec 100644 --- a/src/batchnorm_layer.c +++ b/src/batchnorm_layer.c @@ -2,12 +2,13 @@ #include "blas.h" #include -layer make_batchnorm_layer(int batch, int w, int h, int c) +layer make_batchnorm_layer(int batch, int w, int h, int c, int train) { fprintf(stderr, "Batch Normalization Layer: %d x %d x %d image\n", w,h,c); layer layer = { (LAYER_TYPE)0 }; layer.type = BATCHNORM; layer.batch = batch; + layer.train = train; layer.h = layer.out_h = h; layer.w = layer.out_w = w; layer.c = layer.out_c = c; @@ -42,13 +43,19 @@ layer make_batchnorm_layer(int batch, int w, int h, int c) layer.update_gpu = update_batchnorm_layer_gpu; layer.output_gpu = cuda_make_array(layer.output, h * w * c * batch); - layer.delta_gpu = cuda_make_array(layer.delta, h * w * c * batch); layer.biases_gpu = cuda_make_array(layer.biases, c); - layer.bias_updates_gpu = cuda_make_array(layer.bias_updates, c); - layer.scales_gpu = cuda_make_array(layer.scales, c); - layer.scale_updates_gpu = cuda_make_array(layer.scale_updates, c); + + if (train) { + layer.delta_gpu = cuda_make_array(layer.delta, h * w * c * batch); + + layer.bias_updates_gpu = cuda_make_array(layer.bias_updates, c); + layer.scale_updates_gpu = cuda_make_array(layer.scale_updates, c); + + layer.mean_delta_gpu = cuda_make_array(layer.mean, c); + layer.variance_delta_gpu = cuda_make_array(layer.variance, c); + } layer.mean_gpu = cuda_make_array(layer.mean, c); layer.variance_gpu = cuda_make_array(layer.variance, c); @@ -56,16 +63,18 @@ layer make_batchnorm_layer(int batch, int w, int h, int c) layer.rolling_mean_gpu = cuda_make_array(layer.mean, c); layer.rolling_variance_gpu = cuda_make_array(layer.variance, c); - layer.mean_delta_gpu = cuda_make_array(layer.mean, c); - layer.variance_delta_gpu = cuda_make_array(layer.variance, c); + if (train) { + layer.x_gpu = cuda_make_array(layer.output, layer.batch*layer.outputs); +#ifndef CUDNN + layer.x_norm_gpu = cuda_make_array(layer.output, layer.batch*layer.outputs); +#endif // not CUDNN + } - layer.x_gpu = cuda_make_array(layer.output, layer.batch*layer.outputs); - layer.x_norm_gpu = cuda_make_array(layer.output, layer.batch*layer.outputs); #ifdef CUDNN - cudnnCreateTensorDescriptor(&layer.normTensorDesc); - cudnnCreateTensorDescriptor(&layer.normDstTensorDesc); - cudnnSetTensor4dDescriptor(layer.normDstTensorDesc, CUDNN_TENSOR_NCHW, CUDNN_DATA_FLOAT, layer.batch, layer.out_c, layer.out_h, layer.out_w); - cudnnSetTensor4dDescriptor(layer.normTensorDesc, CUDNN_TENSOR_NCHW, CUDNN_DATA_FLOAT, 1, layer.out_c, 1, 1); + CHECK_CUDNN(cudnnCreateTensorDescriptor(&layer.normTensorDesc)); + CHECK_CUDNN(cudnnCreateTensorDescriptor(&layer.normDstTensorDesc)); + CHECK_CUDNN(cudnnSetTensor4dDescriptor(layer.normDstTensorDesc, CUDNN_TENSOR_NCHW, CUDNN_DATA_FLOAT, layer.batch, layer.out_c, layer.out_h, layer.out_w)); + CHECK_CUDNN(cudnnSetTensor4dDescriptor(layer.normTensorDesc, CUDNN_TENSOR_NCHW, CUDNN_DATA_FLOAT, 1, layer.out_c, 1, 1)); #endif #endif return layer; @@ -129,9 +138,40 @@ void normalize_delta_cpu(float *x, float *mean, float *variance, float *mean_del } } -void resize_batchnorm_layer(layer *layer, int w, int h) +void resize_batchnorm_layer(layer *l, int w, int h) { - fprintf(stderr, "Not implemented\n"); + l->out_h = l->h = h; + l->out_w = l->w = w; + l->outputs = l->inputs = h*w*l->c; + + const int output_size = l->outputs * l->batch; + + l->output = (float*)realloc(l->output, output_size * sizeof(float)); + l->delta = (float*)realloc(l->delta, output_size * sizeof(float)); + +#ifdef GPU + cuda_free(l->output_gpu); + l->output_gpu = cuda_make_array(l->output, output_size); + + if (l->train) { + cuda_free(l->delta_gpu); + l->delta_gpu = cuda_make_array(l->delta, output_size); + + cuda_free(l->x_gpu); + l->x_gpu = cuda_make_array(l->output, output_size); +#ifndef CUDNN + cuda_free(l->x_norm_gpu); + l->x_norm_gpu = cuda_make_array(l->output, output_size); +#endif // not CUDNN + } + + +#ifdef CUDNN + CHECK_CUDNN(cudnnDestroyTensorDescriptor(l->normDstTensorDesc)); + CHECK_CUDNN(cudnnCreateTensorDescriptor(&l->normDstTensorDesc)); + CHECK_CUDNN(cudnnSetTensor4dDescriptor(l->normDstTensorDesc, CUDNN_TENSOR_NCHW, CUDNN_DATA_FLOAT, l->batch, l->out_c, l->out_h, l->out_w)); +#endif // CUDNN +#endif // GPU } void forward_batchnorm_layer(layer l, network_state state) @@ -157,6 +197,7 @@ void forward_batchnorm_layer(layer l, network_state state) normalize_cpu(l.output, l.rolling_mean, l.rolling_variance, l.batch, l.out_c, l.out_h*l.out_w); } scale_bias(l.output, l.scales, l.batch, l.out_c, l.out_h*l.out_w); + add_bias(l.output, l.biases, l.batch, l.out_c, l.out_w*l.out_h); } void backward_batchnorm_layer(const layer l, network_state state) @@ -188,12 +229,14 @@ void update_batchnorm_layer(layer l, int batch, float learning_rate, float momen void pull_batchnorm_layer(layer l) { + cuda_pull_array(l.biases_gpu, l.biases, l.c); cuda_pull_array(l.scales_gpu, l.scales, l.c); cuda_pull_array(l.rolling_mean_gpu, l.rolling_mean, l.c); cuda_pull_array(l.rolling_variance_gpu, l.rolling_variance, l.c); } void push_batchnorm_layer(layer l) { + cuda_push_array(l.biases_gpu, l.biases, l.c); cuda_push_array(l.scales_gpu, l.scales, l.c); cuda_push_array(l.rolling_mean_gpu, l.rolling_mean, l.c); cuda_push_array(l.rolling_variance_gpu, l.rolling_variance, l.c); diff --git a/src/batchnorm_layer.h b/src/batchnorm_layer.h index 55c7d62c..f5137465 100644 --- a/src/batchnorm_layer.h +++ b/src/batchnorm_layer.h @@ -8,11 +8,13 @@ #ifdef __cplusplus extern "C" { #endif -layer make_batchnorm_layer(int batch, int w, int h, int c); +layer make_batchnorm_layer(int batch, int w, int h, int c, int train); void forward_batchnorm_layer(layer l, network_state state); void backward_batchnorm_layer(layer l, network_state state); void update_batchnorm_layer(layer l, int batch, float learning_rate, float momentum, float decay); +void resize_batchnorm_layer(layer *l, int w, int h); + #ifdef GPU void forward_batchnorm_layer_gpu(layer l, network_state state); void backward_batchnorm_layer_gpu(layer l, network_state state); diff --git a/src/blas.h b/src/blas.h index 8477bb41..72c4ae9e 100644 --- a/src/blas.h +++ b/src/blas.h @@ -37,6 +37,7 @@ void mean_cpu(float *x, int batch, int filters, int spatial, float *mean); void variance_cpu(float *x, float *mean, int batch, int filters, int spatial, float *variance); void normalize_cpu(float *x, float *mean, float *variance, int batch, int filters, int spatial); +void add_bias(float *output, float *biases, int batch, int n, int size); void scale_bias(float *output, float *scales, int batch, int n, int size); void backward_scale_cpu(float *x_norm, float *delta, int batch, int n, int size, float *scale_updates); void mean_delta_cpu(float *delta, float *variance, int batch, int filters, int spatial, float *mean_delta); diff --git a/src/network.c b/src/network.c index 4fae07d6..0a2117a8 100644 --- a/src/network.c +++ b/src/network.c @@ -527,6 +527,8 @@ int resize_network(network *net, int w, int h) resize_maxpool_layer(&l, w, h); }else if (l.type == LOCAL_AVGPOOL) { resize_maxpool_layer(&l, w, h); + }else if (l.type == BATCHNORM) { + resize_batchnorm_layer(&l, w, h); }else if(l.type == REGION){ resize_region_layer(&l, w, h); }else if (l.type == YOLO) { @@ -1079,7 +1081,6 @@ void fuse_conv_batchnorm(network net) int f; for (f = 0; f < l->n; ++f) { - //l->biases[f] = l->biases[f] - (double)l->scales[f] * l->rolling_mean[f] / (sqrt((double)l->rolling_variance[f]) + .000001f); l->biases[f] = l->biases[f] - (double)l->scales[f] * l->rolling_mean[f] / (sqrt((double)l->rolling_variance[f] + .000001)); const size_t filter_size = l->size*l->size*l->c / l->groups; @@ -1087,7 +1088,6 @@ void fuse_conv_batchnorm(network net) for (i = 0; i < filter_size; ++i) { int w_index = f*filter_size + i; - //l->weights[w_index] = (double)l->weights[w_index] * l->scales[f] / (sqrt((double)l->rolling_variance[f]) + .000001f); l->weights[w_index] = (double)l->weights[w_index] * l->scales[f] / (sqrt((double)l->rolling_variance[f] + .000001)); } } diff --git a/src/parser.c b/src/parser.c index 0b14d9d6..8da5189e 100644 --- a/src/parser.c +++ b/src/parser.c @@ -805,7 +805,7 @@ layer parse_normalization(list *options, size_params params) layer parse_batchnorm(list *options, size_params params) { - layer l = make_batchnorm_layer(params.batch, params.w, params.h, params.c); + layer l = make_batchnorm_layer(params.batch, params.w, params.h, params.c, params.train); return l; } @@ -1507,6 +1507,7 @@ void save_batchnorm_weights(layer l, FILE *fp) pull_batchnorm_layer(l); } #endif + fwrite(l.biases, sizeof(float), l.c, fp); fwrite(l.scales, sizeof(float), l.c, fp); fwrite(l.rolling_mean, sizeof(float), l.c, fp); fwrite(l.rolling_variance, sizeof(float), l.c, fp); @@ -1652,6 +1653,7 @@ void load_connected_weights(layer l, FILE *fp, int transpose) void load_batchnorm_weights(layer l, FILE *fp) { + fread(l.biases, sizeof(float), l.c, fp); fread(l.scales, sizeof(float), l.c, fp); fread(l.rolling_mean, sizeof(float), l.c, fp); fread(l.rolling_variance, sizeof(float), l.c, fp);