diff --git a/src/batchnorm_layer.c b/src/batchnorm_layer.c index bb363512..a025fc4a 100644 --- a/src/batchnorm_layer.c +++ b/src/batchnorm_layer.c @@ -11,11 +11,15 @@ layer make_batchnorm_layer(int batch, int w, int h, int c) layer.h = layer.out_h = h; layer.w = layer.out_w = w; layer.c = layer.out_c = c; + layer.n = layer.c; layer.output = (float*)calloc(h * w * c * batch, sizeof(float)); layer.delta = (float*)calloc(h * w * c * batch, sizeof(float)); layer.inputs = w*h*c; layer.outputs = layer.inputs; + layer.biases = (float*)calloc(c, sizeof(float)); + layer.bias_updates = (float*)calloc(c, sizeof(float)); + layer.scales = (float*)calloc(c, sizeof(float)); layer.scale_updates = (float*)calloc(c, sizeof(float)); int i; @@ -31,13 +35,18 @@ layer make_batchnorm_layer(int batch, int w, int h, int c) layer.forward = forward_batchnorm_layer; layer.backward = backward_batchnorm_layer; + layer.update = update_batchnorm_layer; #ifdef GPU layer.forward_gpu = forward_batchnorm_layer_gpu; layer.backward_gpu = backward_batchnorm_layer_gpu; + 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); @@ -162,6 +171,19 @@ void backward_batchnorm_layer(const layer l, network_state state) if(l.type == BATCHNORM) copy_cpu(l.outputs*l.batch, l.delta, 1, state.delta, 1); } +void update_batchnorm_layer(layer l, int batch, float learning_rate, float momentum, float decay) +{ + //int size = l.nweights; + axpy_cpu(l.c, learning_rate / batch, l.bias_updates, 1, l.biases, 1); + scal_cpu(l.c, momentum, l.bias_updates, 1); + + axpy_cpu(l.c, learning_rate / batch, l.scale_updates, 1, l.scales, 1); + scal_cpu(l.c, momentum, l.scale_updates, 1); +} + + + + #ifdef GPU void pull_batchnorm_layer(layer l) @@ -213,7 +235,7 @@ void forward_batchnorm_layer_gpu(layer l, network_state state) fix_nan_and_inf(l.rolling_mean_gpu, l.n); fix_nan_and_inf(l.rolling_variance_gpu, l.n); } -#else +#else // CUDNN fast_mean_gpu(l.output_gpu, l.batch, l.out_c, l.out_h*l.out_w, l.mean_gpu); fast_variance_gpu(l.output_gpu, l.mean_gpu, l.batch, l.out_c, l.out_h*l.out_w, l.variance_gpu); @@ -228,7 +250,7 @@ void forward_batchnorm_layer_gpu(layer l, network_state state) scale_bias_gpu(l.output_gpu, l.scales_gpu, l.batch, l.out_c, l.out_h*l.out_w); add_bias_gpu(l.output_gpu, l.biases_gpu, l.batch, l.out_c, l.out_w*l.out_h); -#endif +#endif // CUDNN } else { normalize_gpu(l.output_gpu, l.rolling_mean_gpu, l.rolling_variance_gpu, l.batch, l.out_c, l.out_h*l.out_w); @@ -268,7 +290,7 @@ void backward_batchnorm_layer_gpu(layer l, network_state state) l.variance_gpu); // input (should be FP32) simple_copy_ongpu(l.outputs*l.batch, l.output_gpu, l.delta_gpu); //simple_copy_ongpu(l.outputs*l.batch, l.x_norm_gpu, l.delta_gpu); -#else +#else // CUDNN backward_bias_gpu(l.bias_updates_gpu, l.delta_gpu, l.batch, l.out_c, l.out_w*l.out_h); backward_scale_gpu(l.x_norm_gpu, l.delta_gpu, l.batch, l.out_c, l.out_w*l.out_h, l.scale_updates_gpu); @@ -277,7 +299,7 @@ void backward_batchnorm_layer_gpu(layer l, network_state state) fast_mean_delta_gpu(l.delta_gpu, l.variance_gpu, l.batch, l.out_c, l.out_w*l.out_h, l.mean_delta_gpu); fast_variance_delta_gpu(l.x_gpu, l.delta_gpu, l.mean_gpu, l.variance_gpu, l.batch, l.out_c, l.out_w*l.out_h, l.variance_delta_gpu); normalize_delta_gpu(l.x_gpu, l.mean_gpu, l.variance_gpu, l.mean_delta_gpu, l.variance_delta_gpu, l.batch, l.out_c, l.out_w*l.out_h, l.delta_gpu); -#endif +#endif // CUDNN if (l.type == BATCHNORM) simple_copy_ongpu(l.outputs*l.batch, l.delta_gpu, state.delta); //copy_ongpu(l.outputs*l.batch, l.delta_gpu, 1, state.delta, 1); @@ -286,4 +308,19 @@ void backward_batchnorm_layer_gpu(layer l, network_state state) fix_nan_and_inf(l.bias_updates_gpu, l.n); } } -#endif + +void update_batchnorm_layer_gpu(layer l, int batch, float learning_rate_init, float momentum, float decay) +{ + float learning_rate = learning_rate_init*l.learning_rate_scale; + //float momentum = a.momentum; + //float decay = a.decay; + //int batch = a.batch; + + axpy_ongpu(l.c, learning_rate / batch, l.bias_updates_gpu, 1, l.biases_gpu, 1); + scal_ongpu(l.c, momentum, l.bias_updates_gpu, 1); + + axpy_ongpu(l.c, learning_rate / batch, l.scale_updates_gpu, 1, l.scales_gpu, 1); + scal_ongpu(l.c, momentum, l.scale_updates_gpu, 1); +} + +#endif // GPU diff --git a/src/batchnorm_layer.h b/src/batchnorm_layer.h index 1ea38d72..55c7d62c 100644 --- a/src/batchnorm_layer.h +++ b/src/batchnorm_layer.h @@ -11,10 +11,12 @@ extern "C" { layer make_batchnorm_layer(int batch, int w, int h, int c); 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); #ifdef GPU void forward_batchnorm_layer_gpu(layer l, network_state state); void backward_batchnorm_layer_gpu(layer l, network_state state); +void update_batchnorm_layer_gpu(layer l, int batch, float learning_rate_init, float momentum, float decay); void pull_batchnorm_layer(layer l); void push_batchnorm_layer(layer l); #endif diff --git a/src/convolutional_layer.c b/src/convolutional_layer.c index 1db47fc1..cc05b3b0 100644 --- a/src/convolutional_layer.c +++ b/src/convolutional_layer.c @@ -1447,9 +1447,17 @@ void backward_convolutional_layer(convolutional_layer l, network_state state) } } -void update_convolutional_layer(convolutional_layer l, int batch, float learning_rate, float momentum, float decay) +void update_convolutional_layer(convolutional_layer l, int batch, float learning_rate_init, float momentum, float decay) { - //int size = l.nweights; + float learning_rate = learning_rate_init*l.learning_rate_scale; + //float momentum = a.momentum; + //float decay = a.decay; + //int batch = a.batch; + + axpy_cpu(l.nweights, -decay*batch, l.weights, 1, l.weight_updates, 1); + axpy_cpu(l.nweights, learning_rate / batch, l.weight_updates, 1, l.weights, 1); + scal_cpu(l.nweights, momentum, l.weight_updates, 1); + axpy_cpu(l.n, learning_rate / batch, l.bias_updates, 1, l.biases, 1); scal_cpu(l.n, momentum, l.bias_updates, 1); @@ -1457,10 +1465,6 @@ void update_convolutional_layer(convolutional_layer l, int batch, float learning axpy_cpu(l.n, learning_rate / batch, l.scale_updates, 1, l.scales, 1); scal_cpu(l.n, momentum, l.scale_updates, 1); } - - axpy_cpu(l.nweights, -decay*batch, l.weights, 1, l.weight_updates, 1); - axpy_cpu(l.nweights, learning_rate / batch, l.weight_updates, 1, l.weights, 1); - scal_cpu(l.nweights, momentum, l.weight_updates, 1); }