From d8adaf8ea6a31a380f6bf1fe65e88b661d3bb51e Mon Sep 17 00:00:00 2001
From: Joseph Redmon <pjreddie@gmail.com>
Date: Fri, 21 Oct 2016 20:16:43 +0000
Subject: [PATCH] tree stuff

---
 src/softmax_layer.c |  116 +++++++++++++++++++++++++++++++++++++++++++++++++---------
 1 files changed, 98 insertions(+), 18 deletions(-)

diff --git a/src/softmax_layer.c b/src/softmax_layer.c
index 28696b7..2a34cae 100644
--- a/src/softmax_layer.c
+++ b/src/softmax_layer.c
@@ -1,35 +1,115 @@
 #include "softmax_layer.h"
+#include "blas.h"
+#include "cuda.h"
+#include <float.h>
 #include <math.h>
 #include <stdlib.h>
 #include <stdio.h>
+#include <assert.h>
 
-softmax_layer *make_softmax_layer(int inputs)
+softmax_layer make_softmax_layer(int batch, int inputs, int groups)
 {
-    printf("Softmax Layer: %d inputs\n", inputs);
-    softmax_layer *layer = calloc(1, sizeof(softmax_layer));
-    layer->inputs = inputs;
-    layer->output = calloc(inputs, sizeof(double));
-    layer->delta = calloc(inputs, sizeof(double));
-    return layer;
+    assert(inputs%groups == 0);
+    fprintf(stderr, "Softmax Layer: %d inputs\n", inputs);
+    softmax_layer l = {0};
+    l.type = SOFTMAX;
+    l.batch = batch;
+    l.groups = groups;
+    l.inputs = inputs;
+    l.outputs = inputs;
+    l.output = calloc(inputs*batch, sizeof(float));
+    l.delta = calloc(inputs*batch, sizeof(float));
+
+    l.forward = forward_softmax_layer;
+    l.backward = backward_softmax_layer;
+    #ifdef GPU
+    l.forward_gpu = forward_softmax_layer_gpu;
+    l.backward_gpu = backward_softmax_layer_gpu;
+
+    l.output_gpu = cuda_make_array(l.output, inputs*batch); 
+    l.delta_gpu = cuda_make_array(l.delta, inputs*batch); 
+    #endif
+    return l;
 }
 
-void forward_softmax_layer(const softmax_layer layer, double *input)
+void forward_softmax_layer(const softmax_layer l, network_state state)
+{
+    int b;
+    int inputs = l.inputs / l.groups;
+    int batch = l.batch * l.groups;
+    if(l.softmax_tree){
+        for(b = 0; b < batch; ++b){
+            int i;
+            int count = 0;
+            for(i = 0; i < l.softmax_tree->groups; ++i){
+                int group_size = l.softmax_tree->group_size[i];
+                softmax(state.input+b*inputs + count, group_size, l.temperature, l.output+b*inputs + count);
+                count += group_size;
+            }
+        }
+    } else {
+        for(b = 0; b < batch; ++b){
+            softmax(state.input+b*inputs, inputs, l.temperature, l.output+b*inputs);
+        }
+    }
+}
+
+void backward_softmax_layer(const softmax_layer l, network_state state)
 {
     int i;
-    double sum = 0;
-    for(i = 0; i < layer.inputs; ++i){
-        sum += exp(input[i]);
-    }
-    for(i = 0; i < layer.inputs; ++i){
-        layer.output[i] = exp(input[i])/sum;
+    for(i = 0; i < l.inputs*l.batch; ++i){
+        state.delta[i] += l.delta[i];
     }
 }
 
-void backward_softmax_layer(const softmax_layer layer, double *input, double *delta)
+#ifdef GPU
+
+void pull_softmax_layer_output(const softmax_layer layer)
 {
-    int i;
-    for(i = 0; i < layer.inputs; ++i){
-        delta[i] = layer.delta[i];
+    cuda_pull_array(layer.output_gpu, layer.output, layer.inputs*layer.batch);
+}
+
+void forward_softmax_layer_gpu(const softmax_layer l, network_state state)
+{
+    int inputs = l.inputs / l.groups;
+    int batch = l.batch * l.groups;
+    int b;
+    if(l.softmax_tree){
+        if(0){
+            float *buff = calloc(inputs * batch, sizeof(float));
+            cuda_pull_array(state.input, buff, batch * inputs);
+            state.input = buff;
+            forward_softmax_layer(l, state);
+            cuda_push_array(l.output_gpu, l.output, batch*inputs);
+            free(buff);
+        } else {
+            int i;
+            const int nstreams = 32;
+            cudaStream_t streams[nstreams];
+            for (i = 0; i < nstreams; ++i) {
+                cudaStreamCreate(&streams[i]);
+            }
+            for (b = 0; b < batch; ++b) {
+                int i;
+                int count = 0;
+                for (i = 0; i < l.softmax_tree->groups; ++i) {
+                    int group_size = l.softmax_tree->group_size[i];
+                    softmax_gpu(state.input+b*inputs + count, group_size, 1, l.temperature, l.output_gpu+b*inputs + count, streams[(b*l.softmax_tree->groups + i) % nstreams]);
+                    count += group_size;
+                }
+            }
+            for(i = 0; i < nstreams; ++i){
+                cudaStreamDestroy(streams[i]);
+            }
+        }
+    } else {
+        softmax_gpu(state.input, inputs, batch, l.temperature, l.output_gpu, 0);
     }
 }
 
+void backward_softmax_layer_gpu(const softmax_layer layer, network_state state)
+{
+    axpy_ongpu(layer.batch*layer.inputs, 1, layer.delta_gpu, 1, state.delta, 1);
+}
+
+#endif

--
Gitblit v1.10.0