Another fix for [batchnorm] layer

src/batchnorm_layer.c

@@ -2,12 +2,13 @@
 #include "blas.h"
 #include <stdio.h>
 
-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);
+    if (train) {
-    layer.variance_delta_gpu = cuda_make_array(layer.variance, c);
+        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);
+    CHECK_CUDNN(cudnnCreateTensorDescriptor(&layer.normTensorDesc));
-    cudnnCreateTensorDescriptor(&layer.normDstTensorDesc);
+    CHECK_CUDNN(cudnnCreateTensorDescriptor(&layer.normDstTensorDesc));
-    cudnnSetTensor4dDescriptor(layer.normDstTensorDesc, CUDNN_TENSOR_NCHW, CUDNN_DATA_FLOAT, layer.batch, layer.out_c, layer.out_h, layer.out_w);
+    CHECK_CUDNN(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(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);

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);

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);

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));
                     }
                 }

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);