diff --git a/src/darknet.c b/src/darknet.c index fc560559..5cf11f28 100644 --- a/src/darknet.c +++ b/src/darknet.c @@ -117,6 +117,26 @@ void operations(char *cfgfile) ops += 2l * l.n * l.size*l.size*l.c * l.out_h*l.out_w; } else if(l.type == CONNECTED){ ops += 2l * l.inputs * l.outputs; + } else if (l.type == RNN){ + ops += 2l * l.input_layer->inputs * l.input_layer->outputs; + ops += 2l * l.self_layer->inputs * l.self_layer->outputs; + ops += 2l * l.output_layer->inputs * l.output_layer->outputs; + } else if (l.type == GRU){ + ops += 2l * l.uz->inputs * l.uz->outputs; + ops += 2l * l.uh->inputs * l.uh->outputs; + ops += 2l * l.ur->inputs * l.ur->outputs; + ops += 2l * l.wz->inputs * l.wz->outputs; + ops += 2l * l.wh->inputs * l.wh->outputs; + ops += 2l * l.wr->inputs * l.wr->outputs; + } else if (l.type == LSTM){ + ops += 2l * l.uf->inputs * l.uf->outputs; + ops += 2l * l.ui->inputs * l.ui->outputs; + ops += 2l * l.ug->inputs * l.ug->outputs; + ops += 2l * l.uo->inputs * l.uo->outputs; + ops += 2l * l.wf->inputs * l.wf->outputs; + ops += 2l * l.wi->inputs * l.wi->outputs; + ops += 2l * l.wg->inputs * l.wg->outputs; + ops += 2l * l.wo->inputs * l.wo->outputs; } } printf("Floating Point Operations: %ld\n", ops); @@ -220,6 +240,16 @@ void reset_normalize_net(char *cfgfile, char *weightfile, char *outfile) denormalize_connected_layer(*l.state_r_layer); denormalize_connected_layer(*l.state_h_layer); } + if (l.type == LSTM && l.batch_normalize) { + denormalize_connected_layer(*l.wf); + denormalize_connected_layer(*l.wi); + denormalize_connected_layer(*l.wg); + denormalize_connected_layer(*l.wo); + denormalize_connected_layer(*l.uf); + denormalize_connected_layer(*l.ui); + denormalize_connected_layer(*l.ug); + denormalize_connected_layer(*l.uo); + } } save_weights(net, outfile); } @@ -262,6 +292,17 @@ void normalize_net(char *cfgfile, char *weightfile, char *outfile) *l.state_h_layer = normalize_layer(*l.state_h_layer, l.state_h_layer->outputs); net.layers[i].batch_normalize=1; } + if (l.type == LSTM && l.batch_normalize) { + *l.wf = normalize_layer(*l.wf, l.wf->outputs); + *l.wi = normalize_layer(*l.wi, l.wi->outputs); + *l.wg = normalize_layer(*l.wg, l.wg->outputs); + *l.wo = normalize_layer(*l.wo, l.wo->outputs); + *l.uf = normalize_layer(*l.uf, l.uf->outputs); + *l.ui = normalize_layer(*l.ui, l.ui->outputs); + *l.ug = normalize_layer(*l.ug, l.ug->outputs); + *l.uo = normalize_layer(*l.uo, l.uo->outputs); + net.layers[i].batch_normalize=1; + } } save_weights(net, outfile); } @@ -295,6 +336,25 @@ void statistics_net(char *cfgfile, char *weightfile) printf("State H\n"); statistics_connected_layer(*l.state_h_layer); } + if (l.type == LSTM && l.batch_normalize) { + printf("LSTM Layer %d\n", i); + printf("wf\n"); + statistics_connected_layer(*l.wf); + printf("wi\n"); + statistics_connected_layer(*l.wi); + printf("wg\n"); + statistics_connected_layer(*l.wg); + printf("wo\n"); + statistics_connected_layer(*l.wo); + printf("uf\n"); + statistics_connected_layer(*l.uf); + printf("ui\n"); + statistics_connected_layer(*l.ui); + printf("ug\n"); + statistics_connected_layer(*l.ug); + printf("uo\n"); + statistics_connected_layer(*l.uo); + } printf("\n"); } } @@ -332,6 +392,25 @@ void denormalize_net(char *cfgfile, char *weightfile, char *outfile) l.state_h_layer->batch_normalize = 0; net.layers[i].batch_normalize=0; } + if (l.type == GRU && l.batch_normalize) { + denormalize_connected_layer(*l.wf); + denormalize_connected_layer(*l.wi); + denormalize_connected_layer(*l.wg); + denormalize_connected_layer(*l.wo); + denormalize_connected_layer(*l.uf); + denormalize_connected_layer(*l.ui); + denormalize_connected_layer(*l.ug); + denormalize_connected_layer(*l.uo); + l.wf->batch_normalize = 0; + l.wi->batch_normalize = 0; + l.wg->batch_normalize = 0; + l.wo->batch_normalize = 0; + l.uf->batch_normalize = 0; + l.ui->batch_normalize = 0; + l.ug->batch_normalize = 0; + l.uo->batch_normalize = 0; + net.layers[i].batch_normalize=0; + } } save_weights(net, outfile); } diff --git a/src/lstm_layer.c b/src/lstm_layer.c new file mode 100644 index 00000000..e61bf5c9 --- /dev/null +++ b/src/lstm_layer.c @@ -0,0 +1,626 @@ +#include "lstm_layer.h" +#include "connected_layer.h" +#include "utils.h" +#include "cuda.h" +#include "blas.h" +#include "gemm.h" + +#include +#include +#include +#include + +static void increment_layer(layer *l, int steps) +{ + int num = l->outputs*l->batch*steps; + l->output += num; + l->delta += num; + l->x += num; + l->x_norm += num; + +#ifdef GPU + l->output_gpu += num; + l->delta_gpu += num; + l->x_gpu += num; + l->x_norm_gpu += num; +#endif +} + +layer make_lstm_layer(int batch, int inputs, int outputs, int steps, int batch_normalize) +{ + fprintf(stderr, "LSTM Layer: %d inputs, %d outputs\n", inputs, outputs); + batch = batch / steps; + layer l = { 0 }; + l.batch = batch; + l.type = LSTM; + l.steps = steps; + l.inputs = inputs; + + l.uf = malloc(sizeof(layer)); + fprintf(stderr, "\t\t"); + *(l.uf) = make_connected_layer(batch*steps, inputs, outputs, LINEAR, batch_normalize); + l.uf->batch = batch; + + l.ui = malloc(sizeof(layer)); + fprintf(stderr, "\t\t"); + *(l.ui) = make_connected_layer(batch*steps, inputs, outputs, LINEAR, batch_normalize); + l.ui->batch = batch; + + l.ug = malloc(sizeof(layer)); + fprintf(stderr, "\t\t"); + *(l.ug) = make_connected_layer(batch*steps, inputs, outputs, LINEAR, batch_normalize); + l.ug->batch = batch; + + l.uo = malloc(sizeof(layer)); + fprintf(stderr, "\t\t"); + *(l.uo) = make_connected_layer(batch*steps, inputs, outputs, LINEAR, batch_normalize); + l.uo->batch = batch; + + l.wf = malloc(sizeof(layer)); + fprintf(stderr, "\t\t"); + *(l.wf) = make_connected_layer(batch*steps, outputs, outputs, LINEAR, batch_normalize); + l.wf->batch = batch; + + l.wi = malloc(sizeof(layer)); + fprintf(stderr, "\t\t"); + *(l.wi) = make_connected_layer(batch*steps, outputs, outputs, LINEAR, batch_normalize); + l.wi->batch = batch; + + l.wg = malloc(sizeof(layer)); + fprintf(stderr, "\t\t"); + *(l.wg) = make_connected_layer(batch*steps, outputs, outputs, LINEAR, batch_normalize); + l.wg->batch = batch; + + l.wo = malloc(sizeof(layer)); + fprintf(stderr, "\t\t"); + *(l.wo) = make_connected_layer(batch*steps, outputs, outputs, LINEAR, batch_normalize); + l.wo->batch = batch; + + l.batch_normalize = batch_normalize; + l.outputs = outputs; + + l.output = calloc(outputs*batch*steps, sizeof(float)); + l.state = calloc(outputs*batch, sizeof(float)); + + l.forward = forward_lstm_layer; + l.update = update_lstm_layer; + + l.prev_state_cpu = calloc(batch*outputs, sizeof(float)); + l.prev_cell_cpu = calloc(batch*outputs, sizeof(float)); + l.cell_cpu = calloc(batch*outputs*steps, sizeof(float)); + + l.f_cpu = calloc(batch*outputs, sizeof(float)); + l.i_cpu = calloc(batch*outputs, sizeof(float)); + l.g_cpu = calloc(batch*outputs, sizeof(float)); + l.o_cpu = calloc(batch*outputs, sizeof(float)); + l.c_cpu = calloc(batch*outputs, sizeof(float)); + l.h_cpu = calloc(batch*outputs, sizeof(float)); + l.temp_cpu = calloc(batch*outputs, sizeof(float)); + l.temp2_cpu = calloc(batch*outputs, sizeof(float)); + l.temp3_cpu = calloc(batch*outputs, sizeof(float)); + l.dc_cpu = calloc(batch*outputs, sizeof(float)); + l.dh_cpu = calloc(batch*outputs, sizeof(float)); + +#ifdef GPU + l.forward_gpu = forward_lstm_layer_gpu; + l.backward_gpu = backward_lstm_layer_gpu; + l.update_gpu = update_lstm_layer_gpu; + + l.output_gpu = cuda_make_array(0, batch*outputs*steps); + l.delta_gpu = cuda_make_array(0, batch*l.outputs*steps); + + l.prev_state_gpu = cuda_make_array(0, batch*outputs); + l.prev_cell_gpu = cuda_make_array(0, batch*outputs); + l.cell_gpu = cuda_make_array(0, batch*outputs*steps); + + l.f_gpu = cuda_make_array(0, batch*outputs); + l.i_gpu = cuda_make_array(0, batch*outputs); + l.g_gpu = cuda_make_array(0, batch*outputs); + l.o_gpu = cuda_make_array(0, batch*outputs); + l.c_gpu = cuda_make_array(0, batch*outputs); + l.h_gpu = cuda_make_array(0, batch*outputs); + l.temp_gpu = cuda_make_array(0, batch*outputs); + l.temp2_gpu = cuda_make_array(0, batch*outputs); + l.temp3_gpu = cuda_make_array(0, batch*outputs); + l.dc_gpu = cuda_make_array(0, batch*outputs); + l.dh_gpu = cuda_make_array(0, batch*outputs); +#ifdef CUDNN + cudnnSetTensor4dDescriptor(l.wf->dstTensorDesc, CUDNN_TENSOR_NCHW, CUDNN_DATA_FLOAT, batch, l.wf->out_c, l.wf->out_h, l.wf->out_w); + cudnnSetTensor4dDescriptor(l.wi->dstTensorDesc, CUDNN_TENSOR_NCHW, CUDNN_DATA_FLOAT, batch, l.wi->out_c, l.wi->out_h, l.wi->out_w); + cudnnSetTensor4dDescriptor(l.wg->dstTensorDesc, CUDNN_TENSOR_NCHW, CUDNN_DATA_FLOAT, batch, l.wg->out_c, l.wg->out_h, l.wg->out_w); + cudnnSetTensor4dDescriptor(l.wo->dstTensorDesc, CUDNN_TENSOR_NCHW, CUDNN_DATA_FLOAT, batch, l.wo->out_c, l.wo->out_h, l.wo->out_w); + + cudnnSetTensor4dDescriptor(l.uf->dstTensorDesc, CUDNN_TENSOR_NCHW, CUDNN_DATA_FLOAT, batch, l.uf->out_c, l.uf->out_h, l.uf->out_w); + cudnnSetTensor4dDescriptor(l.ui->dstTensorDesc, CUDNN_TENSOR_NCHW, CUDNN_DATA_FLOAT, batch, l.ui->out_c, l.ui->out_h, l.ui->out_w); + cudnnSetTensor4dDescriptor(l.ug->dstTensorDesc, CUDNN_TENSOR_NCHW, CUDNN_DATA_FLOAT, batch, l.ug->out_c, l.ug->out_h, l.ug->out_w); + cudnnSetTensor4dDescriptor(l.uo->dstTensorDesc, CUDNN_TENSOR_NCHW, CUDNN_DATA_FLOAT, batch, l.uo->out_c, l.uo->out_h, l.uo->out_w); +#endif + +#endif + + return l; +} + +void update_lstm_layer(layer l, int batch, float learning_rate, float momentum, float decay) +{ + update_connected_layer(*(l.wf), batch, learning_rate, momentum, decay); + update_connected_layer(*(l.wi), batch, learning_rate, momentum, decay); + update_connected_layer(*(l.wg), batch, learning_rate, momentum, decay); + update_connected_layer(*(l.wo), batch, learning_rate, momentum, decay); + update_connected_layer(*(l.uf), batch, learning_rate, momentum, decay); + update_connected_layer(*(l.ui), batch, learning_rate, momentum, decay); + update_connected_layer(*(l.ug), batch, learning_rate, momentum, decay); + update_connected_layer(*(l.uo), batch, learning_rate, momentum, decay); +} + +void forward_lstm_layer(layer l, network_state state) +{ + network_state s = { 0 }; + s.train = state.train; + int i; + layer wf = *(l.wf); + layer wi = *(l.wi); + layer wg = *(l.wg); + layer wo = *(l.wo); + + layer uf = *(l.uf); + layer ui = *(l.ui); + layer ug = *(l.ug); + layer uo = *(l.uo); + + fill_cpu(l.outputs * l.batch * l.steps, 0, wf.delta, 1); + fill_cpu(l.outputs * l.batch * l.steps, 0, wi.delta, 1); + fill_cpu(l.outputs * l.batch * l.steps, 0, wg.delta, 1); + fill_cpu(l.outputs * l.batch * l.steps, 0, wo.delta, 1); + + fill_cpu(l.outputs * l.batch * l.steps, 0, uf.delta, 1); + fill_cpu(l.outputs * l.batch * l.steps, 0, ui.delta, 1); + fill_cpu(l.outputs * l.batch * l.steps, 0, ug.delta, 1); + fill_cpu(l.outputs * l.batch * l.steps, 0, uo.delta, 1); + if (state.train) { + fill_cpu(l.outputs * l.batch * l.steps, 0, l.delta, 1); + } + + for (i = 0; i < l.steps; ++i) { + s.input = l.h_cpu; + forward_connected_layer(wf, s); + forward_connected_layer(wi, s); + forward_connected_layer(wg, s); + forward_connected_layer(wo, s); + + s.input = state.input; + forward_connected_layer(uf, s); + forward_connected_layer(ui, s); + forward_connected_layer(ug, s); + forward_connected_layer(uo, s); + + copy_cpu(l.outputs*l.batch, wf.output, 1, l.f_cpu, 1); + axpy_cpu(l.outputs*l.batch, 1, uf.output, 1, l.f_cpu, 1); + + copy_cpu(l.outputs*l.batch, wi.output, 1, l.i_cpu, 1); + axpy_cpu(l.outputs*l.batch, 1, ui.output, 1, l.i_cpu, 1); + + copy_cpu(l.outputs*l.batch, wg.output, 1, l.g_cpu, 1); + axpy_cpu(l.outputs*l.batch, 1, ug.output, 1, l.g_cpu, 1); + + copy_cpu(l.outputs*l.batch, wo.output, 1, l.o_cpu, 1); + axpy_cpu(l.outputs*l.batch, 1, uo.output, 1, l.o_cpu, 1); + + activate_array(l.f_cpu, l.outputs*l.batch, LOGISTIC); + activate_array(l.i_cpu, l.outputs*l.batch, LOGISTIC); + activate_array(l.g_cpu, l.outputs*l.batch, TANH); + activate_array(l.o_cpu, l.outputs*l.batch, LOGISTIC); + + copy_cpu(l.outputs*l.batch, l.i_cpu, 1, l.temp_cpu, 1); + mul_cpu(l.outputs*l.batch, l.g_cpu, 1, l.temp_cpu, 1); + mul_cpu(l.outputs*l.batch, l.f_cpu, 1, l.c_cpu, 1); + axpy_cpu(l.outputs*l.batch, 1, l.temp_cpu, 1, l.c_cpu, 1); + + copy_cpu(l.outputs*l.batch, l.c_cpu, 1, l.h_cpu, 1); + activate_array(l.h_cpu, l.outputs*l.batch, TANH); + mul_cpu(l.outputs*l.batch, l.o_cpu, 1, l.h_cpu, 1); + + copy_cpu(l.outputs*l.batch, l.c_cpu, 1, l.cell_cpu, 1); + copy_cpu(l.outputs*l.batch, l.h_cpu, 1, l.output, 1); + + state.input += l.inputs*l.batch; + l.output += l.outputs*l.batch; + l.cell_cpu += l.outputs*l.batch; + + increment_layer(&wf, 1); + increment_layer(&wi, 1); + increment_layer(&wg, 1); + increment_layer(&wo, 1); + + increment_layer(&uf, 1); + increment_layer(&ui, 1); + increment_layer(&ug, 1); + increment_layer(&uo, 1); + } +} + +void backward_lstm_layer(layer l, network_state state) +{ + network_state s = { 0 }; + s.train = state.train; + int i; + layer wf = *(l.wf); + layer wi = *(l.wi); + layer wg = *(l.wg); + layer wo = *(l.wo); + + layer uf = *(l.uf); + layer ui = *(l.ui); + layer ug = *(l.ug); + layer uo = *(l.uo); + + increment_layer(&wf, l.steps - 1); + increment_layer(&wi, l.steps - 1); + increment_layer(&wg, l.steps - 1); + increment_layer(&wo, l.steps - 1); + + increment_layer(&uf, l.steps - 1); + increment_layer(&ui, l.steps - 1); + increment_layer(&ug, l.steps - 1); + increment_layer(&uo, l.steps - 1); + + state.input += l.inputs*l.batch*(l.steps - 1); + if (state.delta) state.delta += l.inputs*l.batch*(l.steps - 1); + + l.output += l.outputs*l.batch*(l.steps - 1); + l.cell_cpu += l.outputs*l.batch*(l.steps - 1); + l.delta += l.outputs*l.batch*(l.steps - 1); + + for (i = l.steps - 1; i >= 0; --i) { + if (i != 0) copy_cpu(l.outputs*l.batch, l.cell_cpu - l.outputs*l.batch, 1, l.prev_cell_cpu, 1); + copy_cpu(l.outputs*l.batch, l.cell_cpu, 1, l.c_cpu, 1); + if (i != 0) copy_cpu(l.outputs*l.batch, l.output - l.outputs*l.batch, 1, l.prev_state_cpu, 1); + copy_cpu(l.outputs*l.batch, l.output, 1, l.h_cpu, 1); + + l.dh_cpu = (i == 0) ? 0 : l.delta - l.outputs*l.batch; + + copy_cpu(l.outputs*l.batch, wf.output, 1, l.f_cpu, 1); + axpy_cpu(l.outputs*l.batch, 1, uf.output, 1, l.f_cpu, 1); + + copy_cpu(l.outputs*l.batch, wi.output, 1, l.i_cpu, 1); + axpy_cpu(l.outputs*l.batch, 1, ui.output, 1, l.i_cpu, 1); + + copy_cpu(l.outputs*l.batch, wg.output, 1, l.g_cpu, 1); + axpy_cpu(l.outputs*l.batch, 1, ug.output, 1, l.g_cpu, 1); + + copy_cpu(l.outputs*l.batch, wo.output, 1, l.o_cpu, 1); + axpy_cpu(l.outputs*l.batch, 1, uo.output, 1, l.o_cpu, 1); + + activate_array(l.f_cpu, l.outputs*l.batch, LOGISTIC); + activate_array(l.i_cpu, l.outputs*l.batch, LOGISTIC); + activate_array(l.g_cpu, l.outputs*l.batch, TANH); + activate_array(l.o_cpu, l.outputs*l.batch, LOGISTIC); + + copy_cpu(l.outputs*l.batch, l.delta, 1, l.temp3_cpu, 1); + + copy_cpu(l.outputs*l.batch, l.c_cpu, 1, l.temp_cpu, 1); + activate_array(l.temp_cpu, l.outputs*l.batch, TANH); + + copy_cpu(l.outputs*l.batch, l.temp3_cpu, 1, l.temp2_cpu, 1); + mul_cpu(l.outputs*l.batch, l.o_cpu, 1, l.temp2_cpu, 1); + + gradient_array(l.temp_cpu, l.outputs*l.batch, TANH, l.temp2_cpu); + axpy_cpu(l.outputs*l.batch, 1, l.dc_cpu, 1, l.temp2_cpu, 1); + + copy_cpu(l.outputs*l.batch, l.c_cpu, 1, l.temp_cpu, 1); + activate_array(l.temp_cpu, l.outputs*l.batch, TANH); + mul_cpu(l.outputs*l.batch, l.temp3_cpu, 1, l.temp_cpu, 1); + gradient_array(l.o_cpu, l.outputs*l.batch, LOGISTIC, l.temp_cpu); + copy_cpu(l.outputs*l.batch, l.temp_cpu, 1, wo.delta, 1); + s.input = l.prev_state_cpu; + s.delta = l.dh_cpu; + backward_connected_layer(wo, s); + + copy_cpu(l.outputs*l.batch, l.temp_cpu, 1, uo.delta, 1); + s.input = state.input; + s.delta = state.delta; + backward_connected_layer(uo, s); + + copy_cpu(l.outputs*l.batch, l.temp2_cpu, 1, l.temp_cpu, 1); + mul_cpu(l.outputs*l.batch, l.i_cpu, 1, l.temp_cpu, 1); + gradient_array(l.g_cpu, l.outputs*l.batch, TANH, l.temp_cpu); + copy_cpu(l.outputs*l.batch, l.temp_cpu, 1, wg.delta, 1); + s.input = l.prev_state_cpu; + s.delta = l.dh_cpu; + backward_connected_layer(wg, s); + + copy_cpu(l.outputs*l.batch, l.temp_cpu, 1, ug.delta, 1); + s.input = state.input; + s.delta = state.delta; + backward_connected_layer(ug, s); + + copy_cpu(l.outputs*l.batch, l.temp2_cpu, 1, l.temp_cpu, 1); + mul_cpu(l.outputs*l.batch, l.g_cpu, 1, l.temp_cpu, 1); + gradient_array(l.i_cpu, l.outputs*l.batch, LOGISTIC, l.temp_cpu); + copy_cpu(l.outputs*l.batch, l.temp_cpu, 1, wi.delta, 1); + s.input = l.prev_state_cpu; + s.delta = l.dh_cpu; + backward_connected_layer(wi, s); + + copy_cpu(l.outputs*l.batch, l.temp_cpu, 1, ui.delta, 1); + s.input = state.input; + s.delta = state.delta; + backward_connected_layer(ui, s); + + copy_cpu(l.outputs*l.batch, l.temp2_cpu, 1, l.temp_cpu, 1); + mul_cpu(l.outputs*l.batch, l.prev_cell_cpu, 1, l.temp_cpu, 1); + gradient_array(l.f_cpu, l.outputs*l.batch, LOGISTIC, l.temp_cpu); + copy_cpu(l.outputs*l.batch, l.temp_cpu, 1, wf.delta, 1); + s.input = l.prev_state_cpu; + s.delta = l.dh_cpu; + backward_connected_layer(wf, s); + + copy_cpu(l.outputs*l.batch, l.temp_cpu, 1, uf.delta, 1); + s.input = state.input; + s.delta = state.delta; + backward_connected_layer(uf, s); + + copy_cpu(l.outputs*l.batch, l.temp2_cpu, 1, l.temp_cpu, 1); + mul_cpu(l.outputs*l.batch, l.f_cpu, 1, l.temp_cpu, 1); + copy_cpu(l.outputs*l.batch, l.temp_cpu, 1, l.dc_cpu, 1); + + state.input -= l.inputs*l.batch; + if (state.delta) state.delta -= l.inputs*l.batch; + l.output -= l.outputs*l.batch; + l.cell_cpu -= l.outputs*l.batch; + l.delta -= l.outputs*l.batch; + + increment_layer(&wf, -1); + increment_layer(&wi, -1); + increment_layer(&wg, -1); + increment_layer(&wo, -1); + + increment_layer(&uf, -1); + increment_layer(&ui, -1); + increment_layer(&ug, -1); + increment_layer(&uo, -1); + } +} + +#ifdef GPU +void update_lstm_layer_gpu(layer l, int batch, float learning_rate, float momentum, float decay) +{ + update_connected_layer_gpu(*(l.wf), batch, learning_rate, momentum, decay); + update_connected_layer_gpu(*(l.wi), batch, learning_rate, momentum, decay); + update_connected_layer_gpu(*(l.wg), batch, learning_rate, momentum, decay); + update_connected_layer_gpu(*(l.wo), batch, learning_rate, momentum, decay); + update_connected_layer_gpu(*(l.uf), batch, learning_rate, momentum, decay); + update_connected_layer_gpu(*(l.ui), batch, learning_rate, momentum, decay); + update_connected_layer_gpu(*(l.ug), batch, learning_rate, momentum, decay); + update_connected_layer_gpu(*(l.uo), batch, learning_rate, momentum, decay); +} + +void forward_lstm_layer_gpu(layer l, network_state state) +{ + network_state s = { 0 }; + s.train = state.train; + int i; + layer wf = *(l.wf); + layer wi = *(l.wi); + layer wg = *(l.wg); + layer wo = *(l.wo); + + layer uf = *(l.uf); + layer ui = *(l.ui); + layer ug = *(l.ug); + layer uo = *(l.uo); + + fill_ongpu(l.outputs * l.batch * l.steps, 0, wf.delta_gpu, 1); + fill_ongpu(l.outputs * l.batch * l.steps, 0, wi.delta_gpu, 1); + fill_ongpu(l.outputs * l.batch * l.steps, 0, wg.delta_gpu, 1); + fill_ongpu(l.outputs * l.batch * l.steps, 0, wo.delta_gpu, 1); + + fill_ongpu(l.outputs * l.batch * l.steps, 0, uf.delta_gpu, 1); + fill_ongpu(l.outputs * l.batch * l.steps, 0, ui.delta_gpu, 1); + fill_ongpu(l.outputs * l.batch * l.steps, 0, ug.delta_gpu, 1); + fill_ongpu(l.outputs * l.batch * l.steps, 0, uo.delta_gpu, 1); + if (state.train) { + fill_ongpu(l.outputs * l.batch * l.steps, 0, l.delta_gpu, 1); + } + + for (i = 0; i < l.steps; ++i) { + s.input = l.state_gpu; + forward_connected_layer_gpu(wf, s); + forward_connected_layer_gpu(wi, s); + forward_connected_layer_gpu(wg, s); + forward_connected_layer_gpu(wo, s); + + s.input = state.input; + forward_connected_layer_gpu(uf, s); + forward_connected_layer_gpu(ui, s); + forward_connected_layer_gpu(ug, s); + forward_connected_layer_gpu(uo, s); + + copy_ongpu(l.outputs*l.batch, wf.output_gpu, 1, l.f_gpu, 1); + axpy_ongpu(l.outputs*l.batch, 1, uf.output_gpu, 1, l.f_gpu, 1); + + copy_ongpu(l.outputs*l.batch, wi.output_gpu, 1, l.i_gpu, 1); + axpy_ongpu(l.outputs*l.batch, 1, ui.output_gpu, 1, l.i_gpu, 1); + + copy_ongpu(l.outputs*l.batch, wg.output_gpu, 1, l.g_gpu, 1); + axpy_ongpu(l.outputs*l.batch, 1, ug.output_gpu, 1, l.g_gpu, 1); + + copy_ongpu(l.outputs*l.batch, wo.output_gpu, 1, l.o_gpu, 1); + axpy_ongpu(l.outputs*l.batch, 1, uo.output_gpu, 1, l.o_gpu, 1); + + activate_array_ongpu(l.f_gpu, l.outputs*l.batch, LOGISTIC); + activate_array_ongpu(l.i_gpu, l.outputs*l.batch, LOGISTIC); + activate_array_ongpu(l.g_gpu, l.outputs*l.batch, TANH); + activate_array_ongpu(l.o_gpu, l.outputs*l.batch, LOGISTIC); + + copy_ongpu(l.outputs*l.batch, l.i_gpu, 1, l.temp_gpu, 1); + mul_ongpu(l.outputs*l.batch, l.g_gpu, 1, l.temp_gpu, 1); + mul_ongpu(l.outputs*l.batch, l.f_gpu, 1, l.c_gpu, 1); + axpy_ongpu(l.outputs*l.batch, 1, l.temp_gpu, 1, l.c_gpu, 1); + + copy_ongpu(l.outputs*l.batch, l.c_gpu, 1, l.h_gpu, 1); + activate_array_ongpu(l.h_gpu, l.outputs*l.batch, TANH); + mul_ongpu(l.outputs*l.batch, l.o_gpu, 1, l.h_gpu, 1); + + copy_ongpu(l.outputs*l.batch, l.c_gpu, 1, l.cell_gpu, 1); + copy_ongpu(l.outputs*l.batch, l.h_gpu, 1, l.output_gpu, 1); + + state.input += l.inputs*l.batch; + l.output_gpu += l.outputs*l.batch; + l.cell_gpu += l.outputs*l.batch; + + increment_layer(&wf, 1); + increment_layer(&wi, 1); + increment_layer(&wg, 1); + increment_layer(&wo, 1); + + increment_layer(&uf, 1); + increment_layer(&ui, 1); + increment_layer(&ug, 1); + increment_layer(&uo, 1); + } +} + +void backward_lstm_layer_gpu(layer l, network_state state) +{ + network_state s = { 0 }; + s.train = state.train; + int i; + layer wf = *(l.wf); + layer wi = *(l.wi); + layer wg = *(l.wg); + layer wo = *(l.wo); + + layer uf = *(l.uf); + layer ui = *(l.ui); + layer ug = *(l.ug); + layer uo = *(l.uo); + + increment_layer(&wf, l.steps - 1); + increment_layer(&wi, l.steps - 1); + increment_layer(&wg, l.steps - 1); + increment_layer(&wo, l.steps - 1); + + increment_layer(&uf, l.steps - 1); + increment_layer(&ui, l.steps - 1); + increment_layer(&ug, l.steps - 1); + increment_layer(&uo, l.steps - 1); + + state.input += l.inputs*l.batch*(l.steps - 1); + if (state.delta) state.delta += l.inputs*l.batch*(l.steps - 1); + + l.output_gpu += l.outputs*l.batch*(l.steps - 1); + l.cell_gpu += l.outputs*l.batch*(l.steps - 1); + l.delta_gpu += l.outputs*l.batch*(l.steps - 1); + + for (i = l.steps - 1; i >= 0; --i) { + if (i != 0) copy_ongpu(l.outputs*l.batch, l.cell_gpu - l.outputs*l.batch, 1, l.prev_cell_gpu, 1); + copy_ongpu(l.outputs*l.batch, l.cell_gpu, 1, l.c_gpu, 1); + if (i != 0) copy_ongpu(l.outputs*l.batch, l.output_gpu - l.outputs*l.batch, 1, l.prev_state_gpu, 1); + copy_ongpu(l.outputs*l.batch, l.output_gpu, 1, l.h_gpu, 1); + + l.dh_gpu = (i == 0) ? 0 : l.delta_gpu - l.outputs*l.batch; + + copy_ongpu(l.outputs*l.batch, wf.output_gpu, 1, l.f_gpu, 1); + axpy_ongpu(l.outputs*l.batch, 1, uf.output_gpu, 1, l.f_gpu, 1); + + copy_ongpu(l.outputs*l.batch, wi.output_gpu, 1, l.i_gpu, 1); + axpy_ongpu(l.outputs*l.batch, 1, ui.output_gpu, 1, l.i_gpu, 1); + + copy_ongpu(l.outputs*l.batch, wg.output_gpu, 1, l.g_gpu, 1); + axpy_ongpu(l.outputs*l.batch, 1, ug.output_gpu, 1, l.g_gpu, 1); + + copy_ongpu(l.outputs*l.batch, wo.output_gpu, 1, l.o_gpu, 1); + axpy_ongpu(l.outputs*l.batch, 1, uo.output_gpu, 1, l.o_gpu, 1); + + activate_array_ongpu(l.f_gpu, l.outputs*l.batch, LOGISTIC); + activate_array_ongpu(l.i_gpu, l.outputs*l.batch, LOGISTIC); + activate_array_ongpu(l.g_gpu, l.outputs*l.batch, TANH); + activate_array_ongpu(l.o_gpu, l.outputs*l.batch, LOGISTIC); + + copy_ongpu(l.outputs*l.batch, l.delta_gpu, 1, l.temp3_gpu, 1); + + copy_ongpu(l.outputs*l.batch, l.c_gpu, 1, l.temp_gpu, 1); + activate_array_ongpu(l.temp_gpu, l.outputs*l.batch, TANH); + + copy_ongpu(l.outputs*l.batch, l.temp3_gpu, 1, l.temp2_gpu, 1); + mul_ongpu(l.outputs*l.batch, l.o_gpu, 1, l.temp2_gpu, 1); + + gradient_array_ongpu(l.temp_gpu, l.outputs*l.batch, TANH, l.temp2_gpu); + axpy_ongpu(l.outputs*l.batch, 1, l.dc_gpu, 1, l.temp2_gpu, 1); + + copy_ongpu(l.outputs*l.batch, l.c_gpu, 1, l.temp_gpu, 1); + activate_array_ongpu(l.temp_gpu, l.outputs*l.batch, TANH); + mul_ongpu(l.outputs*l.batch, l.temp3_gpu, 1, l.temp_gpu, 1); + gradient_array_ongpu(l.o_gpu, l.outputs*l.batch, LOGISTIC, l.temp_gpu); + copy_ongpu(l.outputs*l.batch, l.temp_gpu, 1, wo.delta_gpu, 1); + s.input = l.prev_state_gpu; + s.delta = l.dh_gpu; + backward_connected_layer_gpu(wo, s); + + copy_ongpu(l.outputs*l.batch, l.temp_gpu, 1, uo.delta_gpu, 1); + s.input = state.input; + s.delta = state.delta; + backward_connected_layer_gpu(uo, s); + + copy_ongpu(l.outputs*l.batch, l.temp2_gpu, 1, l.temp_gpu, 1); + mul_ongpu(l.outputs*l.batch, l.i_gpu, 1, l.temp_gpu, 1); + gradient_array_ongpu(l.g_gpu, l.outputs*l.batch, TANH, l.temp_gpu); + copy_ongpu(l.outputs*l.batch, l.temp_gpu, 1, wg.delta_gpu, 1); + s.input = l.prev_state_gpu; + s.delta = l.dh_gpu; + backward_connected_layer_gpu(wg, s); + + copy_ongpu(l.outputs*l.batch, l.temp_gpu, 1, ug.delta_gpu, 1); + s.input = state.input; + s.delta = state.delta; + backward_connected_layer_gpu(ug, s); + + copy_ongpu(l.outputs*l.batch, l.temp2_gpu, 1, l.temp_gpu, 1); + mul_ongpu(l.outputs*l.batch, l.g_gpu, 1, l.temp_gpu, 1); + gradient_array_ongpu(l.i_gpu, l.outputs*l.batch, LOGISTIC, l.temp_gpu); + copy_ongpu(l.outputs*l.batch, l.temp_gpu, 1, wi.delta_gpu, 1); + s.input = l.prev_state_gpu; + s.delta = l.dh_gpu; + backward_connected_layer_gpu(wi, s); + + copy_ongpu(l.outputs*l.batch, l.temp_gpu, 1, ui.delta_gpu, 1); + s.input = state.input; + s.delta = state.delta; + backward_connected_layer_gpu(ui, s); + + copy_ongpu(l.outputs*l.batch, l.temp2_gpu, 1, l.temp_gpu, 1); + mul_ongpu(l.outputs*l.batch, l.prev_cell_gpu, 1, l.temp_gpu, 1); + gradient_array_ongpu(l.f_gpu, l.outputs*l.batch, LOGISTIC, l.temp_gpu); + copy_ongpu(l.outputs*l.batch, l.temp_gpu, 1, wf.delta_gpu, 1); + s.input = l.prev_state_gpu; + s.delta = l.dh_gpu; + backward_connected_layer_gpu(wf, s); + + copy_ongpu(l.outputs*l.batch, l.temp_gpu, 1, uf.delta_gpu, 1); + s.input = state.input; + s.delta = state.delta; + backward_connected_layer_gpu(uf, s); + + copy_ongpu(l.outputs*l.batch, l.temp2_gpu, 1, l.temp_gpu, 1); + mul_ongpu(l.outputs*l.batch, l.f_gpu, 1, l.temp_gpu, 1); + copy_ongpu(l.outputs*l.batch, l.temp_gpu, 1, l.dc_gpu, 1); + + state.input -= l.inputs*l.batch; + if (state.delta) state.delta -= l.inputs*l.batch; + l.output_gpu -= l.outputs*l.batch; + l.cell_gpu -= l.outputs*l.batch; + l.delta_gpu -= l.outputs*l.batch; + + increment_layer(&wf, -1); + increment_layer(&wi, -1); + increment_layer(&wg, -1); + increment_layer(&wo, -1); + + increment_layer(&uf, -1); + increment_layer(&ui, -1); + increment_layer(&ug, -1); + increment_layer(&uo, -1); + } +} +#endif diff --git a/src/lstm_layer.h b/src/lstm_layer.h new file mode 100644 index 00000000..ad37c96f --- /dev/null +++ b/src/lstm_layer.h @@ -0,0 +1,20 @@ +#ifndef LSTM_LAYER_H +#define LSTM_LAYER_H + +#include "activations.h" +#include "layer.h" +#include "network.h" +#define USET + +layer make_lstm_layer(int batch, int inputs, int outputs, int steps, int batch_normalize); + +void forward_lstm_layer(layer l, network_state state); +void update_lstm_layer(layer l, int batch, float learning_rate, float momentum, float decay); + +#ifdef GPU +void forward_lstm_layer_gpu(layer l, network_state state); +void backward_lstm_layer_gpu(layer l, network_state state); +void update_lstm_layer_gpu(layer l, int batch, float learning_rate, float momentum, float decay); + +#endif +#endif diff --git a/src/network.c b/src/network.c index 465764f2..79686aa0 100644 --- a/src/network.c +++ b/src/network.c @@ -140,6 +140,8 @@ char *get_layer_string(LAYER_TYPE a) return "rnn"; case GRU: return "gru"; + case LSTM: + return "lstm"; case CRNN: return "crnn"; case MAXPOOL: diff --git a/src/parser.c b/src/parser.c index 1a4c8b82..3a51ab8b 100644 --- a/src/parser.c +++ b/src/parser.c @@ -18,6 +18,7 @@ #include "gru_layer.h" #include "list.h" #include "local_layer.h" +#include "lstm_layer.h" #include "maxpool_layer.h" #include "normalization_layer.h" #include "option_list.h" @@ -58,6 +59,7 @@ LAYER_TYPE string_to_layer_type(char * type) || strcmp(type, "[network]")==0) return NETWORK; if (strcmp(type, "[crnn]")==0) return CRNN; if (strcmp(type, "[gru]")==0) return GRU; + if (strcmp(type, "[lstm]")==0) return LSTM; if (strcmp(type, "[rnn]")==0) return RNN; if (strcmp(type, "[conn]")==0 || strcmp(type, "[connected]")==0) return CONNECTED; @@ -219,6 +221,16 @@ layer parse_gru(list *options, size_params params) return l; } +layer parse_lstm(list *options, size_params params) +{ + int output = option_find_int(options, "output",1); + int batch_normalize = option_find_int_quiet(options, "batch_normalize", 0); + + layer l = make_lstm_layer(params.batch, params.inputs, output, params.time_steps, batch_normalize); + + return l; +} + connected_layer parse_connected(list *options, size_params params) { int output = option_find_int(options, "output",1); @@ -755,6 +767,8 @@ network parse_network_cfg_custom(char *filename, int batch) l = parse_rnn(options, params); }else if(lt == GRU){ l = parse_gru(options, params); + }else if(lt == LSTM){ + l = parse_lstm(options, params); }else if(lt == CRNN){ l = parse_crnn(options, params); }else if(lt == CONNECTED){ @@ -1025,6 +1039,15 @@ void save_weights_upto(network net, char *filename, int cutoff) save_connected_weights(*(l.state_z_layer), fp); save_connected_weights(*(l.state_r_layer), fp); save_connected_weights(*(l.state_h_layer), fp); + } if(l.type == LSTM){ + save_connected_weights(*(l.wf), fp); + save_connected_weights(*(l.wi), fp); + save_connected_weights(*(l.wg), fp); + save_connected_weights(*(l.wo), fp); + save_connected_weights(*(l.uf), fp); + save_connected_weights(*(l.ui), fp); + save_connected_weights(*(l.ug), fp); + save_connected_weights(*(l.uo), fp); } if(l.type == CRNN){ save_convolutional_weights(*(l.input_layer), fp); save_convolutional_weights(*(l.self_layer), fp); @@ -1236,6 +1259,16 @@ void load_weights_upto(network *net, char *filename, int cutoff) load_connected_weights(*(l.state_r_layer), fp, transpose); load_connected_weights(*(l.state_h_layer), fp, transpose); } + if(l.type == LSTM){ + load_connected_weights(*(l.wf), fp, transpose); + load_connected_weights(*(l.wi), fp, transpose); + load_connected_weights(*(l.wg), fp, transpose); + load_connected_weights(*(l.wo), fp, transpose); + load_connected_weights(*(l.uf), fp, transpose); + load_connected_weights(*(l.ui), fp, transpose); + load_connected_weights(*(l.ug), fp, transpose); + load_connected_weights(*(l.uo), fp, transpose); + } if(l.type == LOCAL){ int locations = l.out_w*l.out_h; int size = l.size*l.size*l.c*l.n*locations; @@ -1281,4 +1314,4 @@ network *load_network(char *cfg, char *weights, int clear) } if (clear) (*net->seen) = 0; return net; -} \ No newline at end of file +}