Remove test_training_sgd_many_iters()

Author: Maksymilian Graczyk

tests/test_lenet.py

@@ -201,7 +201,7 @@ def loss_grad(layer, b):
                     pytorch_layer.bias.grad, 1e-11)
 
 
-def run_training(net_arguments, mnist, iterations):
+def run_training(net_arguments, mnist):
     mnist_train, _ = mnist
 
     net, pytorch_net, layers = net_arguments
@@ -216,54 +216,44 @@ def run_training(net_arguments, mnist, iterations):
                       pytorch_net.fc1, pytorch_net.fc2, pytorch_net.fc3]
     devito_layers = [layers[0], layers[2], layers[5], layers[6], layers[7]]
 
-    epsilon = 1.26e-11
-
-    for i, data in enumerate(mnist_train, 0):
-        images, labels = data
-
-        def loss_grad(layer, b):
-            gradients = []
+    images, labels = iter(mnist_train).next()
 
-            for j in range(10):
-                result = layer.result.data[j, b]
-                if j == labels[b]:
-                    result -= 1
-                gradients.append(result)
+    def loss_grad(layer, b):
+        gradients = []
 
-            return gradients
+        for j in range(10):
+            result = layer.result.data[j, b]
+            if j == labels[b]:
+                result -= 1
+            gradients.append(result)
 
-        images = images.double()
+        return gradients
 
-        outputs = net.forward(images.numpy())
+    images = images.double()
 
-        pytorch_optimizer.zero_grad()
-        pytorch_outputs = pytorch_net(images)
+    outputs = net.forward(images.numpy())
 
-        compare(outputs, nn.Softmax(dim=1)(pytorch_outputs),
-                1e-12 + i * epsilon)
+    pytorch_optimizer.zero_grad()
+    pytorch_outputs = pytorch_net(images)
 
-        net.backward(loss_grad, optimizer)
+    compare(outputs, nn.Softmax(dim=1)(pytorch_outputs),
+            1e-12)
 
-        pytorch_loss = criterion(pytorch_outputs, labels)
-        pytorch_loss.backward()
-        pytorch_optimizer.step()
+    net.backward(loss_grad, optimizer)
 
-        for j in range(len(pytorch_layers) - 1, -1, -1):
-            pytorch_layer = pytorch_layers[j]
-            devito_layer = devito_layers[j]
+    pytorch_loss = criterion(pytorch_outputs, labels)
+    pytorch_loss.backward()
+    pytorch_optimizer.step()
 
-            compare(devito_layer.kernel.data, pytorch_layer.weight,
-                    1e-12 + i * epsilon)
-            compare(devito_layer.bias.data, pytorch_layer.bias,
-                    1e-12 + i * epsilon)
+    for j in range(len(pytorch_layers) - 1, -1, -1):
+        pytorch_layer = pytorch_layers[j]
+        devito_layer = devito_layers[j]
 
-        if i == iterations - 1:
-            break
+        compare(devito_layer.kernel.data, pytorch_layer.weight,
+                1e-12)
+        compare(devito_layer.bias.data, pytorch_layer.bias,
+                1e-12)
 
 
 def test_training_sgd(net_arguments, mnist):
-    run_training(net_arguments, mnist, 1)
-
-
-def test_training_sgd_many_iters(net_arguments, mnist):
-    run_training(net_arguments, mnist, 20)
+    run_training(net_arguments, mnist)