From 0305fb4d99cf1efc7d4aa4d2ee2d65d54500d437 Mon Sep 17 00:00:00 2001
From: Joseph Redmon <pjreddie@gmail.com>
Date: Thu, 26 Nov 2015 19:48:01 +0000
Subject: [PATCH] Some changes

---
 src/network.c |   12 ++++++++++++
 1 files changed, 12 insertions(+), 0 deletions(-)

diff --git a/src/network.c b/src/network.c
index 6c7461d..d9585c4 100644
--- a/src/network.c
+++ b/src/network.c
@@ -26,6 +26,17 @@
     return batch_num;
 }
 
+void reset_momentum(network net)
+{
+    if (net.momentum == 0) return;
+    net.learning_rate = 0;
+    net.momentum = 0;
+    net.decay = 0;
+    #ifdef GPU
+        if(gpu_index >= 0) update_network_gpu(net);
+    #endif
+}
+
 float get_current_rate(network net)
 {
     int batch_num = get_current_batch(net);
@@ -41,6 +52,7 @@
             for(i = 0; i < net.num_steps; ++i){
                 if(net.steps[i] > batch_num) return rate;
                 rate *= net.scales[i];
+                if(net.steps[i] > batch_num - 1) reset_momentum(net);
             }
             return rate;
         case EXP:

--
Gitblit v1.10.0