lib/nn.rb in nn-1.5 vs lib/nn.rb in nn-1.6

@@ -1,10 +1,10 @@
 require "numo/narray"
 require "json"
 
 class NN
-  VERSION = "1.5"
+  VERSION = "1.6"
 
   include Numo
 
   attr_accessor :weights
   attr_accessor :biases
@@ -20,11 +20,11 @@
 
   def initialize(num_nodes,
                  learning_rate: 0.01,
                  batch_size: 1,
                  activation: %i(relu identity),
-                 momentum: 0,
+                 momentum: 0.9,
                  weight_decay: 0,
                  use_dropout: false,
                  dropout_ratio: 0.5,
                  use_batch_norm: false)
     SFloat.srand(rand(2 ** 64))
@@ -62,33 +62,41 @@
       nn.betas = json["betas"].map{|beta| SFloat.cast(beta)}
     end
     nn
   end
 
-  def train(x_train, y_train, epoch,
-            save_dir: nil, save_interval: 1, test: nil, border: nil, tolerance: 0.5, &block)
+  def train(x_train, y_train, epochs,
+            learning_rate_decay: 0,
+            save_dir: nil,
+            save_interval: 1,
+            test: nil,
+            border: nil,
+            tolerance: 0.5,
+            &block)
     num_train_data = x_train.is_a?(SFloat) ? x_train.shape[0] : x_train.length
-    (epoch * num_train_data / @batch_size).times do |count|
-      loss = learn(x_train, y_train, &block)
-      if loss.nan?
-        puts "loss is nan"
-        break
-      end
-      if (count + 1) % (num_train_data / @batch_size) == 0
-        now_epoch = (count + 1) / (num_train_data / @batch_size)
-        if save_dir && now_epoch % save_interval == 0
-          save("#{save_dir}/epoch#{now_epoch}.json")
+    (1..epochs).each do |epoch|
+      loss = nil
+      (num_train_data.to_f / @batch_size).ceil.times do
+        loss = learn(x_train, y_train, &block)
+        if loss.nan?
+          puts "loss is nan"
+          break
         end
-        msg = "epoch #{now_epoch}/#{epoch} loss: #{loss}"
-        if test
-          acc = accurate(*test, tolerance, &block)
-          puts "#{msg} accurate: #{acc}"
-          break if border && acc >= border
-        else
-          puts msg
-        end
       end
+      if save_dir && epoch % save_interval == 0
+        save("#{save_dir}/epoch#{epoch}.json")
+      end
+      msg = "epoch #{epoch}/#{epochs} loss: #{loss}"
+      if test
+        acc = accurate(*test, tolerance, &block)
+        puts "#{msg} accurate: #{acc}"
+        break if border && acc >= border
+      else
+        puts msg
+      end
+      @learning_rate -= learning_rate_decay
+      @learning_rate = 1e-7 if @learning_rate < 1e-7
     end
   end
 
   def test(x_test, y_test, tolerance = 0.5, &block)
     acc = accurate(x_test, y_test, tolerance, &block)
@@ -97,10 +105,10 @@
   end
 
   def accurate(x_test, y_test, tolerance = 0.5, &block)
     correct = 0
     num_test_data = x_test.is_a?(SFloat) ? x_test.shape[0] : x_test.length
-    (num_test_data / @batch_size).times do |i|
+    (num_test_data.to_f / @batch_size).ceil.times do |i|
       x = SFloat.zeros(@batch_size, @num_nodes.first)
       y = SFloat.zeros(@batch_size, @num_nodes.last)
       @batch_size.times do |j|
         k = i * @batch_size + j
         if x_test.is_a?(SFloat)