add test and fix cifar10 normalization values (#438)

Author: mikowals

Datasets/CIFAR10/CIFAR10.swift

@@ -98,9 +98,30 @@ func loadCIFARFile(named name: String, in directory: URL, normalizing: Bool = tr
     // Transpose from the CIFAR-provided N(CHW) to TF's default NHWC.
     var imageTensor = Tensor<Float>(images.transposed(permutation: [0, 2, 3, 1]))
 
+    // The value of mean and std were calculated with the following Swift code:
+    // ```
+    // import TensorFlow
+    // import Datasets
+    // import Foundation
+    // let urlString = "https://storage.googleapis.com/s4tf-hosted-binaries/datasets/CIFAR10/cifar-10-binary.tar.gz"
+    // let cifar = CIFAR10(batchSize: 50000,
+    //                     remoteBinaryArchiveLocation: URL(string: urlString)!,
+    //                     normalizing: false)
+    // for batch in cifar.training.sequenced() {
+    //     let images = Tensor<Double>(batch.first) / 255.0
+    //     let mom = images.moments(squeezingAxes: [0,1,2])
+    //     print("mean: \(mom.mean) std: \(sqrt(mom.variance))")
+    // }
+    // ```
     if normalizing {
-        let mean = Tensor<Float>([0.485, 0.456, 0.406])
-        let std = Tensor<Float>([0.229, 0.224, 0.225])
+        let mean = Tensor<Float>(
+                [0.4913996898,
+                 0.4821584196,
+                 0.4465309242])
+        let std = Tensor<Float>(
+                [0.2470322324,
+                 0.2434851280,
+                 0.2615878417])
         imageTensor = ((imageTensor / 255.0) - mean) / std
     }
 

Tests/DatasetsTests/CIFAR10/CIFAR10Tests.swift

@@ -25,10 +25,34 @@ final class CIFAR10Tests: XCTestCase {
         }
         XCTAssertEqual(totalCount, 50000)
     }
+    
+    func testNormalizeCIFAR10() {
+        let dataset = CIFAR10(
+            batchSize: 50000,
+            remoteBinaryArchiveLocation:
+                URL(
+                    string:
+                        "https://storage.googleapis.com/s4tf-hosted-binaries/datasets/CIFAR10/cifar-10-binary.tar.gz"
+                )!, normalizing: true
+        )
+        
+        let targetMean = Tensor<Double>([0, 0, 0])
+        let targetStd = Tensor<Double>([1, 1, 1])
+        for batch in dataset.training.sequenced() {
+            let images = Tensor<Double>(batch.first)
+            let mean = images.mean(squeezingAxes: [0, 1, 2])
+            let std = images.standardDeviation(squeezingAxes: [0, 1, 2])
+            XCTAssertTrue(targetMean.isAlmostEqual(to: mean,
+                                                   tolerance: 1e-6))
+            XCTAssertTrue(targetStd.isAlmostEqual(to: std,
+                                                  tolerance: 1e-5))
+        }
+    }
 }
 
 extension CIFAR10Tests {
     static var allTests = [
         ("testCreateCIFAR10", testCreateCIFAR10),
+        ("testNormalizeCIFAR10", testNormalizeCIFAR10),
     ]
-}
+}