From 81f7fc2c7bbddb79b5c87b5ab88fedf98f2b9963 Mon Sep 17 00:00:00 2001
From: AlexeyAB <alexeyab84@gmail.com>
Date: Mon, 21 Jan 2019 15:32:57 +0300
Subject: [PATCH] Fixed network resize memory allocation

---
 src/detector.c |  3 ++-
 src/network.c  | 16 ++++++++++++++++
 2 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/src/detector.c b/src/detector.c
index 04572189..747cf422 100644
--- a/src/detector.c
+++ b/src/detector.c
@@ -181,7 +181,8 @@ void train_detector(char *datacfg, char *cfgfile, char *weightfile, int *gpus, i
             int dim_w = roundl(random_val*init_w / 32 + 1) * 32;
             int dim_h = roundl(random_val*init_h / 32 + 1) * 32;
 
-            if (get_current_batch(net) == 0) {
+            // at the beginning
+            if (avg_loss < 0) {
                 dim_w = roundl(1.4*init_w / 32 + 1) * 32;
                 dim_h = roundl(1.4*init_h / 32 + 1) * 32;
             }
diff --git a/src/network.c b/src/network.c
index 465764f2..0017b884 100644
--- a/src/network.c
+++ b/src/network.c
@@ -400,6 +400,12 @@ int resize_network(network *net, int w, int h)
             cuda_free(*net->truth_gpu);
             *net->truth_gpu = 0;
         }
+
+        if (net->input_state_gpu) cuda_free(net->input_state_gpu);
+        if (net->input_pinned_cpu) {
+            if (net->input_pinned_cpu_flag) cudaFreeHost(net->input_pinned_cpu);
+            else free(net->input_pinned_cpu);
+        }
     }
 #endif
     int i;
@@ -448,14 +454,24 @@ int resize_network(network *net, int w, int h)
         h = l.out_h;
         if(l.type == AVGPOOL) break;
     }
+    const int size = get_network_input_size(*net) * net->batch;
 #ifdef GPU
     if(gpu_index >= 0){
         printf(" try to allocate workspace = %zu * sizeof(float), ", workspace_size / sizeof(float) + 1);
         net->workspace = cuda_make_array(0, workspace_size/sizeof(float) + 1);
+        net->input_state_gpu = cuda_make_array(0, size);
+        if (cudaSuccess == cudaHostAlloc(&net->input_pinned_cpu, size * sizeof(float), cudaHostRegisterMapped))
+            net->input_pinned_cpu_flag = 1;
+        else {
+            net->input_pinned_cpu = calloc(size, sizeof(float));
+            net->input_pinned_cpu_flag = 0;
+        }
         printf(" CUDA allocate done! \n");
     }else {
         free(net->workspace);
         net->workspace = calloc(1, workspace_size);
+        if(!net->input_pinned_cpu_flag)
+            net->input_pinned_cpu = realloc(net->input_pinned_cpu, size * sizeof(float));
     }
 #else
     free(net->workspace);