mnist_example2.rb in ruby-dnn-0.2.0

- old
+ new

@@ -10,22 +10,22 @@
 MNIST = DNN::MNIST
 
 x_train, y_train = MNIST.load_train
 x_test, y_test = MNIST.load_test
 
-x_train = SFloat.cast(x_train).reshape(x_train.shape[0], 1, 28, 28)
-x_test = SFloat.cast(x_test).reshape(x_test.shape[0], 1, 28, 28)
+x_train = SFloat.cast(x_train).reshape(x_train.shape[0], 28, 28, 1)
+x_test = SFloat.cast(x_test).reshape(x_test.shape[0], 28, 28, 1)
 
 x_train /= 255
 x_test /= 255
 
 y_train = DNN::Util.to_categorical(y_train, 10)
 y_test = DNN::Util.to_categorical(y_test, 10)
 
 model = Model.new
 
-model << InputLayer.new([1, 28, 28])
+model << InputLayer.new([28, 28, 1])
 
 model << Conv2D.new(16, 5, 5)
 model << BatchNormalization.new
 model << ReLU.new
 
@@ -45,8 +45,6 @@
 model << Dense.new(10)
 model << SoftmaxWithLoss.new
 
 model.compile(Adam.new)
 
-model.train(x_train, y_train, 10, batch_size: 100) do
-  model.test(x_test, y_test)
-end
+model.train(x_train, y_train, 10, batch_size: 100, test: [x_test, y_test])